#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Jan 18 10:04:17 EST 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 cmip6_model_members_in_common import cmip6_model_members_in_common
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
ocean_only = True if 'ocean' in sys.argv else False
tc_grids = True if 'tc' in sys.argv else False
normlev = 1000 # None
years = slice(1982,2014)
years_reanalysis = None# slice(2001,2020)
#years_reanalysis = slice(2001,2020)
#years = slice(2001,2014)
#years_reanalysis = slice(2001,2014)
#cmip6
#ifile = 'ta.historical.annualmean.tropmean.1982-2014.nc'
ifile = '/tigress/wenchang/data/cmip6/variables/ta/historical/wy_analysis/ta.historical.annualmean.tropmean.1982-2014.nc'
if ocean_only:
    ifile = ifile.replace('.nc', '.oceanonly.nc')
elif tc_grids:
    pass #todo
daname, expname = ifile.split('.')[0:2]

ofile = __file__.replace('.py', f'_cmip6.nc')
if ocean_only:
    ofile = ofile.replace('.nc', '.oceanonly.nc')
elif tc_grids:
    ofile = ofile.replace('.nc', '.oceanonly.nc') #todo#ofile = ofile.replace('.nc', '.tcgrids.nc')
if os.path.exists(ofile):
    rg_pi = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    da = xr.open_dataarray(ifile)
    rg = da.linregress.on(da.year)
    models = [s.split('_')[0] for s in rg.ens.values]
    rg = rg.assign_coords(model=('ens', models))
    #subset ens
    model_members = cmip6_model_members_in_common(('PI', 'historical'), ('PI', 'hist-nat'), ('PI', 'hist-GHG'), yearspan=(1982,2014))
    rg_pi = rg.sel(ens=model_members) #subset of ens for which PI is available
    models = [s.split('_')[0] for s in model_members]
    #models = xr.DataArray(models, dims='ens')
    rg_pi = rg_pi.assign_coords(model=('ens', models))
    #save
    rg_pi.to_netcdf(ofile)
    print('[saved]:', ofile)

#era5
ofile = __file__.replace('.py', f'_era5.nc')
if ocean_only:
    ofile = ofile.replace('.nc', '.oceanonly.nc')
elif tc_grids:
    ofile = ofile.replace('.nc', '.tcgrids.nc')
if os.path.exists(ofile):
    rg_era5 = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    ifile = '/tigress/wenchang/data/era5/raw/monthly/plevels/temperature/wyanalysis/era5.temperature.annualmean.tropmean.1979-2020.nc'
    if ocean_only:
        ifile = ifile.replace('.nc', '.oceanonly.nc')
        da = xr.open_dataarray(ifile)
    elif tc_grids:
        ifile = '/tigress/wenchang/data/era5/analysis_wy/TCI_v202201/wyanalysis/era5.yearly.temperature.t.global.FMA-ASO.1979-2020.nc'
        da = xr.open_dataset(ifile)['t']
    else:
        da = xr.open_dataarray(ifile)
    if years_reanalysis is not None:
        da = da.sel(year=years_reanalysis)
    else:
        da = da.sel(year=years)
    rg_era5 = da.linregress.on(da.year)
    #save
    rg_era5.to_netcdf(ofile)
    print('[saved]:', ofile)

#merra2
ofile = __file__.replace('.py', f'_merra2.nc')
if ocean_only:
    ofile = ofile.replace('.nc', '.oceanonly.nc')
elif tc_grids:
    ofile = ofile.replace('.nc', '.tcgrids.nc')
if os.path.exists(ofile):
    rg_merra2 = xr.open_dataset(ofile)
    print('[loaded]:', ofile)
else:
    ifile = '/tigress/wenchang/data/merra2/raw/monthly/instM_3d_asm_Np/T/wyanalysis/merra2.T.annualmean.tropmean.1980-2020.nc'
    if ocean_only:
        ifile = ifile.replace('.nc', '.oceanonly.nc')
        da = xr.open_dataarray(ifile)
    elif tc_grids:
        ifile = '/tigress/wenchang/data/merra2/analysis/TCI/wyanalysis/merra2.yearly.Ta.T.global.FMA-ASO.1980-2020.nc'
        da = xr.open_dataset(ifile)['T']
    else:
        da = xr.open_dataarray(ifile)
    if years_reanalysis is not None:
        da = da.sel(year=years_reanalysis)
    else:
        da = da.sel(year=years)
    rg_merra2 = da.linregress.on(da.year)
    #save
    rg_merra2.to_netcdf(ofile)
    print('[saved]:', ofile)

 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig, ax = plt.subplots(figsize=(5, 5))
    yscale = 'linear'
    plevs = slice(1000,50)
    #normalized_by_surface = True and False
    """
    #all ens
    da = rg.slope.pipe(lambda x: x*10).assign_attrs(units='K/decade')
    da = da.assign_coords(plev=da.plev/100)
    da.plev.attrs['units'] = 'hPa'
    for ens in da.ens.values:
        da.sel(ens=ens).plot(color='C0', ls='-', lw=0.5,  alpha=0.1, y='plev', yincrease=False, yscale=yscale)
    da.mean('ens').plot(color='C0', ls='--', lw=2,  y='plev', yincrease=False, yscale=yscale, label='all ens, simple ens-mean')
    da.groupby(da.model).mean('ens').mean('model').plot(color='C0', ls='-', lw=2,  y='plev', yincrease=False, yscale=yscale, label='all ens, wgt ens-mean')
    """

    #subset of ens
    da = rg_pi.slope.pipe(lambda x: x*10).assign_attrs(units='K/decade')
    da = da.assign_coords(plev=da.plev/100)
    da.plev.attrs['units'] = 'hPa'
    da = da.sel(plev=plevs)
    #if normalized_by_surface:
    if normlev is not None:
        da = da/da.sel(plev=normlev)
    for ens in da.ens.values:
        da.sel(ens=ens).plot(color='C1', ls='-', lw=0.5,  alpha=0.1, y='plev', yincrease=False, yscale=yscale)
    #da.mean('ens').plot(color='C1', ls='--', lw=2,  y='plev', yincrease=False, yscale=yscale, label=f'CMIP6 {expname}, simple ens-mean')
    da.groupby(da.model).mean('ens').mean('model').plot(color='C1', ls='--', lw=2,  y='plev', yincrease=False, yscale=yscale, label=f'CMIP6 {expname}, wgt ens-mean')

    #era5
    if years_reanalysis is not None:
        label = f'ERA5 {years_reanalysis.start}-{years_reanalysis.stop}'
    else:
        label = 'ERA5'
    da = rg_era5.slope.pipe(lambda x: x*10).assign_attrs(long_name=f'Ta linear trend over {years.start}-{years.stop}', units='K/decade')
    da.level.attrs['units'] = 'hPa'
    da = da.isel(level=slice(-1, None, -1)).sel(level=plevs)
    #if normalized_by_surface:
    if normlev is not None:
        da = da/da.sel(level=normlev)
    da.plot(y='level', yscale=yscale, yincrease=False, color='k', lw=2, ls='-', label=label)

    #merra2
    if years_reanalysis is not None:
        label = f'MERRA2 {years_reanalysis.start}-{years_reanalysis.stop}'
    else:
        label = 'MERRA2'
    long_name = f'Annual tropical Ta linear trend over {years.start}-{years.stop}'
    units = 'K/decade'
    da = rg_merra2.slope.pipe(lambda x: x*10).assign_attrs(long_name=long_name, units=units)
    da.lev.attrs['units'] = 'hPa'
    da = da.sel(lev=plevs)
    #if normalized_by_surface:
    if normlev is not None:
        da = da/da.sel(lev=normlev)
        long_name = f'Normalized annual tropical Ta linear trend over {years.start}-{years.stop}'
        units = None 
        da = da.assign_attrs(long_name=long_name, units=units)
    da.plot(y='lev', yscale=yscale, yincrease=False, color='gray', lw=2, ls='-', label=label)

    plt.legend()
    #plt.xlabel(f'CMIP6 {expname} annual tropical {daname} linear trend over 1982-2014 [K/decade]')
    #if normalized_by_surface:
    if normlev is not None:
        plt.axvline(1, color='gray', ls='--')
    else:
        plt.axvline(0, color='gray', ls='--')
    
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'.png')
        #if normalized_by_surface:
        if normlev is not None:
            figname = figname.replace('.png', f'_norm{normlev}.png')
        if years_reanalysis is not None:
            figname = figname.replace('.png', f'_RA{years_reanalysis.start}-{years_reanalysis.stop}.png')
        if ocean_only:
            figname = figname.replace('.png', '_oceanonly.png')
        elif tc_grids:
            figname = figname.replace('.png', '_tcgrids.png')
        if 'pdf' in sys.argv:
            figname = figname.replace('.png', '.pdf')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()