#!/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'
#ifile = 'ts_historical_545ens_1980-2014.nc'
daname, expname = ifile.split('_')[0:2]
#years = slice('1958', '2014')
year0, year1 = ifile.split('.')[-2].split('_')[-1].split('-')
years = slice(year0, year1)
ds = xr.open_dataset(ifile).load()
n = ds.model.groupby(ds.model).count()
#subset of models with 10+ ens
counts_bigens = n.sel(model=(n>=10)) #only models with 10+ ens
models_bigens = n.model.sel(model=(n>=10)) #only models with 10+ ens
ds = ds.sel(model_member=(ds.model.isin(models_bigens)))
ds = ds.groupby(ds.model).mean('model_member')#ens mean with each model
#print(counts_bigens, models_bigens)
#print(ds)
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    da = ds.slope.mean('model').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'{expname} {daname} linear trend {ds.model.size}-model wgt-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'_{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'{expname} {daname} linear trend over {years.start}-{years.stop}')
    g = da.plot.contourf(levels=33, vmax=3.2, center=0, extend='both', col='model', col_wrap=3,
        cbar_kwargs={'orientation': 'vertical', 'shrink': 0.6, 'pad': 0.2}, figsize=(10,6))
    for ii,ax in enumerate(g.axes.flat[0:ds.model.size]):
        mapplot(xticks=range(0,360,90), yticks=range(-30,31,15), fill_continents=True, ax=ax)
        cs = da.isel(model=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('')
        #ax.text(0.01, 0.98, f'{ii+1}', ha='left', va='top', fontsize='x-small', color='lightgray', transform=ax.transAxes)
        ax.set_title(f'{da.model[ii].item()} (n_ens={counts_bigens[ii].item()})', fontsize='medium')
    plt.tight_layout(h_pad=1, w_pad=0, rect=[0,0,0.85,1])
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:] and ds.model.size>1:
        figname = __file__.replace('.py', f'_{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()