#!/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 geoxarray
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
do_anomaly = [True, False][-1]
singleEnsClim = [True, False][-1]
yearspan_clim = slice(2010,2020)
def get_time_series(da, years_clim=None):
    if years_clim is None: years_clim = slice(1980,1990)
    dims_clim = ['year', 'ens'] if 'ens' in da.dims else 'year' #clim average dims
    if singleEnsClim: dims_clim = 'year'
    if do_anomaly:
        return da.geo.fldmean() \
            .groupby('time.year').mean('time') \
            .pipe(lambda x: x - x.sel(year=years_clim).mean(dims_clim))
    else:
        return da.geo.fldmean() \
            .groupby('time.year').mean('time')  \
            .pipe(lambda x: x - 273.15) #K -> C

def do_linear_trend(da):
    """get linear trend for yearly time series"""
    rg = da.linregress.on(da.year, dim='year')
    dyear = da.year[-1].item() - da.year[0].item() #time interval in terms of years
    return rg.slope.pipe(lambda x: x*dyear).assign_attrs(units=f'degC per {dyear} years', long_name='Tas linear trend')


#merra2
ifile = 'merra2.T2M.1980-2020.westUS.nc'
da = xr.open_dataarray(ifile).pipe(get_time_series, years_clim=yearspan_clim)
da_merra2 = da

#era5
ifile = 'era5.2m_temperature.1980-2020.westUS.nc'
da = xr.open_dataarray(ifile).pipe(get_time_series, years_clim=yearspan_clim)
da_era5 = da

#am2.5c360
ifiles = """
t_ref_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_10ens_1980-2019_westUS.nc
t_ref_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_extend2020_10ens_2020-2020_westUS.nc
""".split()
#print(ifiles); sys.exit()
#da = xr.open_mfdataset(ifiles)['rh_ref'].pipe(get_time_series)
ifile = 't_ref_AM2.5C360_amipHadISSTlong_chancorr_tigercpu_intelmpi_18_1080PE_10ens_1980-2020_westUS.nc' #concat file
da = xr.open_dataarray(ifile).pipe(get_time_series, years_clim=yearspan_clim)
da_amip = da

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    import seaborn as sns
    fig, ax = plt.subplots()

    #merra2
    da = da_merra2
    da.plot(color='gray', label='MERRA2')
    #era5
    da = da_era5
    da.plot(color='k', label='ERA5')

    da = da_amip
    da.plot(hue='ens', color='C0', lw=1, add_legend=False, alpha=0.2, ls='-')
    da.mean('ens').plot(label='AM2.5C360 ens mean', ls='-', color='C0')

    ax.legend()
    ylabel = f'Tas anomaly from {yearspan_clim.start}-{yearspan_clim.stop}clim [degC]' if do_anomaly else 'Tas [degC]'
    ax.set_ylabel(ylabel)

    
    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if not do_anomaly: figname = figname.replace('.png', '__absTas.png')
        if do_anomaly: 
            if singleEnsClim: figname = figname.replace('.png', '__singleEnsClim.png')
            if yearspan_clim.start != 1980 or yearspan_clim.stop != 1990: figname = figname.replace('.png', f'__clim{yearspan_clim.start}-{yearspan_clim.stop}.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)

    #linear trend histogram
    fig, ax = plt.subplots()
    #am2.5c360
    da = da_amip.pipe(do_linear_trend)
    sns.kdeplot(data=da, color='C0', ax=ax)
    sns.rugplot(data=da, color='C0', ax=ax)
    ax.axvline(da.mean('ens'), color='C0', label='AM2.5C360 ens mean')
    #da.plot.hist(label='AM2.5C360', bins=bins)
    #merra2
    da = da_merra2.pipe(do_linear_trend)
    ax.axvline(da, color='gray', label='MERRA2')
    #era5
    da = da_era5.pipe(do_linear_trend)
    ax.axvline(da, color='k', label='ERA5')

    ax.legend()
    ax.set_xlabel(f'{da.attrs["long_name"]} ({da.attrs["units"]})')

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'__hist_lineartrend.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()