#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun  8 10:50:03 EDT 2022
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    tt = Timer('start ' + ' '.join(sys.argv))
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import xlinregress
from misc.landmask import flagland
maskland = lambda x: x.where(flagland(x)<0.5)
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ifile = 'ts_historical_545ens_1958-2014.nc'
daname, expname = ifile.split('_')[0:2]
model = 'CanESM5'
years = slice('1958', '2014')
ds = xr.open_dataset(ifile).load()
ds = ds.sel(model_member=(ds.model==model))

 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    da = ds.slope.mean('model_member').pipe(maskland).sel(lat=slice(-30,30)) \
        .pipe(lambda x: x*100).assign_attrs(units='K/century', long_name='linear trend')
    da.plot.contourf(levels=33, vmax=3.2, center=0, ax=ax, extend='both')
    cs = da.plot.contour(levels=[1,], colors='gray')
    ax.clabel(cs)
    mapplot(xticks=np.arange(0,360,60), fill_continents=True)
    title = f'{model} {expname} {daname} linear trend {ds.model_member.size}-ens mean over {years.start}-{years.stop}'
    ax.set_title(title)
    ax.set_xlabel('')
    ax.set_ylabel('')

    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{model}_{expname}_tdensmean_{years.start}-{years.stop}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    
    #fig, ax = plt.subplots(figsize=(10,6))
    da = ds.slope.pipe(maskland).sel(lat=slice(-30,30)).pipe(lambda x: x*100) \
        .assign_attrs(units='K/century', long_name=f'{model} {expname} {daname} linear trend over {years.start}-{years.stop}')
    g = da.plot.contourf(levels=33, vmax=3.2, center=0, extend='both', col='model_member', col_wrap=10,
        cbar_kwargs={'orientation': 'horizontal', 'shrink': 0.6}, figsize=(11,5))
    for ii,ax in enumerate(g.axes.flat[0:ds.model_member.size]):
        mapplot(xticks=[], yticks=[], fill_continents=True, ax=ax)
        cs = da.isel(model_member=ii).plot.contour(levels=[1,], colors='gray', ax=ax, linewidths=0.5)
        ax.clabel(cs, fontsize='x-small')
        ax.set_xlabel('')
        ax.set_ylabel('')
        ax.set_title('')
    plt.tight_layout(pad=0, rect=[0,0.25,1,1])
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:] and ds.model_member.size>1:
        figname = __file__.replace('.py', f'_{model}_{expname}_{years.start}-{years.stop}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv:
        pass
    else:
        plt.show()