#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Apr  7 19:14:37 EDT 2024
if __name__ == '__main__':
    import sys,os
    from misc.timer import Timer
    tt = Timer(f'[{os.getcwd()}] 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
wython = '/tigress/wenchang/wython'
if wython not in sys.path: sys.path.append(wython); print('added to python path:', wython)
#from misc import get_kws_from_argv
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
yearspan = slice(1990,2020)
nyears = yearspan.stop - yearspan.start
def get_lintrend(da):
    da_yearly = da.groupby('time.year').mean('time').sel(year=yearspan)
    ds = da_yearly.linregress.on(da_yearly.year, dim='year')
    return ds

#era5
ifile = 'era5.2m_relative_humidity.1980-2020.westUS.nc'
ds = xr.open_dataarray(ifile).pipe(get_lintrend)
ds_era5 = ds

#am2.5c360
ifiles = """
rh_ref_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_10ens_1980-2019_westUS.nc
rh_ref_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_extend2020_10ens_2020-2020_westUS.nc
""".split()
#print(ifiles); sys.exit()
ifile = 'rh_ref_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_10ens_1980-2020_westUS.nc' #concat file
ds = xr.open_dataarray(ifile).pipe(get_lintrend)
ds_amip = ds

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    fig, axes = plt.subplots(1, 2, figsize=(8, 4))
    func = lambda x: (x*nyears).assign_attrs(units=f'% per {nyears} years', long_name='linear trend')
    cmap = 'BrBG'
    ax = axes[0]
    da = ds_era5.slope.pipe(func)
    da.plot(ax=ax, cmap=cmap, levels=np.arange(-15, 15.1, 2), extend='both')
    ax.set_title('ERA5')

    ax = axes[1]
    da = ds_amip.slope.mean('ens').pipe(func)
    da.plot(ax=ax, cmap=cmap, center=0, levels=np.arange(-15, 15.1, 2), extend='both')
    ax.set_title('AM2.5C360 ens mean')

    for ax in axes: 
        mapplot(ax=ax)
        ax.set_xlabel('')
        ax.set_ylabel('')
    fig.suptitle(f'RH2m linear trend over {yearspan.start}-{yearspan.stop}')

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if yearspan.start != 1980 or yearspan.stop != 2020: figname = figname.replace('.png', f'__{yearspan.start}-{yearspan.stop}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 
    #ens members
    da = ds_amip.slope.pipe(func)
    figsize = 9,8
    fig,axes = plt.subplots(3,4, figsize=figsize, sharex=True, sharey=True, dpi=100)
    for ii,ax in zip(da.ens.values, axes.flat):
        da_ = da.sel(ens=ii)
        im = da_.plot(ax=ax, levels=np.arange(-15, 15.1, 2), cmap='BrBG', extend='both', add_colorbar=False)
        mapplot(ax=ax)
        ax.set_xlabel('')
        ax.set_ylabel('')
        #ax.set_xlim(360-125, 360-95)
        #ax.set_ylim(25, 50)
    for ax in axes.flat[-2:]: ax.set_visible(False)
    fig.colorbar(im, ax=axes.flat[-2:], orientation='horizontal', pad=-0.6, shrink=0.6, label=da.attrs['units'])
    fig.suptitle(f'AM2.5C360 RH2m linear trend over {yearspan.start}-{yearspan.stop}')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__members.png')
        if yearspan.start != 1980 or yearspan.stop != 2020: figname = figname.replace('.png', f'__{yearspan.start}-{yearspan.stop}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
 

    tt.check(f'**Done**')
    print()
    if 'notshowfig' in sys.argv or 'n' in sys.argv:
        pass
    else:
        if 'plt' in globals(): plt.show()