#!/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()
"""
#subset of models with 10+ ens
n = ds.model.groupby(ds.model).count()
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)
"""
#first ens member of each model
ds = ds.groupby(ds.model).first()
da_lintrend = ds.slope.pipe(maskland).sel(lat=slice(-5,5)).mean('lat')
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot, xticks2lon
    fig, ax = plt.subplots(figsize=(8,4))
    func_units = lambda x: (x*100).assign_attrs(units='K/century', long_name='linear trend')
    da = da_lintrend.pipe(func_units)
    for ii,m in enumerate(da.model.values):
        label = f'{da.model.size} CMIP6 models' if ii==0 else None
        da.sel(model=m).plot(ax=ax, label=label, lw=1, alpha=0.3, color='C0', ls='-')
    da.mean('model').plot(ax=ax, label='mmm', lw=2, color='k', ls=':')

    #obs
    #ERSST5
    xr.open_dataset(f'fig_map_lintrend_ersst5_{years.start}-{years.stop}.nc').slope.sel(lat=slice(-5,5)).load().mean('lat') \
        .pipe(func_units) \
        .plot(ax=ax, label='ERSST5', lw=2, color='k', ls='--')
    #HadISST1
    xr.open_dataset(f'fig_map_lintrend_hadisst1_{years.start}-{years.stop}.nc').slope.sel(lat=slice(5,-5)).load().mean('lat') \
        .pipe(func_units) \
        .plot(ax=ax, label='HadISST1', lw=2, color='k', ls='-')
    #HadISST1c
    xr.open_dataset(f'fig_map_lintrend_hadisst1c_{years.start}-{years.stop}.nc').slope.sel(lat=slice(-5,5)).load().mean('lat') \
        .pipe(func_units) \
        .plot(ax=ax, label='HadISST1c', lw=2, color='r', ls='--')



    ax.legend(ncol=1, bbox_to_anchor=[1,1.04], loc='upper left')
    xticks2lon(range(0,360,60), ax=ax)
    title = f'SST linear trend from CMIP6({da.model.size}models) and Obs. {years.start}-{years.stop}, 5S-5N'
    ax.set_title(title)
    ax.set_xlabel('')
    
    #title = f'{expname} {daname} linear trend {ds.model.size}-model wgt-mean over {years.start}-{years.stop}'
    #ax.set_title(title)
    #ax.set_ylabel('')

    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[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()