#!/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 = 'slp_FLOR_HistRCP45_tigercpu_intelmpi_18_576PE_5ens_1860-2100_monthly.nc'
ifile = 'slp_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_10ens_1871-2019_monthly.nc'
daname, model, expname = ifile.split('_')[:3]
years = slice('1958', '2018')
da = xr.open_dataarray(ifile).sel(time=years, lat=slice(-30,30)).load().groupby('time.year').mean('time')
da_em = da.mean('ens')
rg = da.linregress.on(da.year, dim='year')
rg_em = da_em.linregress.on(da_em.year, dim='year')


 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, ax = plt.subplots()
    da = rg_em.slope.pipe(maskland).pipe(lambda x: x*100).assign_attrs(units='hPa/century', long_name='linear trend')
    da.plot.contourf(levels=45, vmax=2.2, ax=ax, cmap='PRGn_r')
    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 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}_ensmean_{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 = rg.slope.pipe(maskland).pipe(lambda x: x*100) \
        .assign_attrs(units='hPa/century', long_name=f'{model} {expname} {daname} linear trend over {years.start}-{years.stop}')
    orientation = 'horizontal' if model=='FLOR' else 'vertical'
    g = da.plot.contourf(levels=45, vmax=2.2, extend='both', col='ens', col_wrap=3, cmap='PRGn_r', figsize=(10,6),
        cbar_kwargs={'orientation': orientation, 'shrink': 0.6})
    for ii,ax in zip(da.ens.values, g.axes.flat):
        mapplot(xticks=np.arange(0,360,60), fill_continents=True, ax=ax)
        cs = da.sel(ens=ii).plot.contour(levels=[1,], colors='gray', ax=ax)
        ax.clabel(cs)
        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}_{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()