#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sat Dec 18 21:22:10 EST 2021
if __name__ == '__main__':
    import sys
    from misc.timer import Timer
    s = ' '
    tt = Timer(f'start {s.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
from misc.landmask import flagland
import xlinregress
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
"""
#cal the weights used to merge NH and SH to global
da = xr.open_dataarray('/tigress/wenchang/data/cmip6/variables/ts/hist-GHG/wy_regrid_all_members/ts.hist-GHG.ACCESS-CM2.r1i1p1f1.1850-2020.nc')
is_ocean = flagland(da) < 0.1
NH = is_ocean.sel(lat=slice(10,30), lon=slice(40,360-20)) 
NH = NH.weighted(np.cos(np.deg2rad(NH.lat))).sum(['lat', 'lon'])
SH = is_ocean.sel(lat=slice(-30,-10), lon=slice(30,360-150)) 
SH = SH.weighted(np.cos(np.deg2rad(SH.lat))).sum(['lat', 'lon'])
wgtSH = 1/(NH/SH + 1)
wgtNH = 1 - wgtSH
wgts = (wgtNH.item(), wgtSH.item())
print('weights of NH and SH: ', wgts)
"""

#savedata = True if len(sys.argv)>1 and 'savedata' in sys.argv[1:] else False
def get_obs_data(dsname, obs='ERA5'):
    if obs == 'ERA5':
        idir = '/tigress/wenchang/data/era5/analysis_wy/TCI_v202201/wyanalysis/excludeIndianOcean'
        year0 = 1979
    elif obs == 'MERRA2':
        idir = '/tigress/wenchang/data/merra2/analysis/TCI/wyanalysis/excludeIndianOcean'
        year0 = 1980
    print(f'{dsname} {obs}')
    #dsname = 'PI'
    print(dsname)
    if dsname == 'PI':
        ifile = os.path.join(idir, f'{obs.lower()}.yearly.PI.vmax.global.FMA-ASO.{year0}-2020.noIO.nc')
        daname = 'vmax'
        dsname_new = 'PI'
        long_name = 'Potential Intensity'
        units = 'm/s'
    elif dsname == 'SST':
        ifile = os.path.join(idir, f'{obs.lower()}.yearly.sst.sst.global.FMA-ASO.{year0}-2020.noIO.nc')
        daname = 'sst'
        dsname_new = 'SST'
        long_name = 'SST'
        units = 'K'
    elif dsname == 'RH':
        ifile = os.path.join(idir, f'{obs.lower()}.yearly.RH.RH.global.FMA-ASO.{year0}-2020.noIO.nc')
        daname = 'RH'
        dsname_new = 'RH'
        long_name = 'Relative Humidity'
        units = '%'
    elif dsname == 'Vshear':
        ifile = os.path.join(idir, f'{obs.lower()}.yearly.Vshear.Vshear.global.FMA-ASO.{year0}-2020.noIO.nc')
        daname = 'Vshear'
        dsname_new = 'Vshear'
        long_name = 'Vertical Wind Shear'
        units = 'm/s'
    trend_units = f'{units}/year'
    len_season = 3
    #
    da = xr.open_dataset(ifile)[daname]
    da = da.sel(year=slice(1982, 2014))
    da = da.assign_attrs(long_name=long_name, units=units)
    lintrend = da.linregress.on(da.year)
    lintrend.slope.attrs['units'] = trend_units
    
    #
    ds = lintrend
    ds['timeseries'] = da
    
    #save data
    ofile = f'{dsname_new}.global.tcseason{len_season}.{obs}.wy.noIO.nc'
    if os.path.exists(ofile):
        print('[exists]:', ofile)
    else:
        ds.to_netcdf(ofile)
        print('[saved]:', ofile)
 
if __name__ == '__main__':
    #from wyconfig import * #my plot settings
    dsnames = ('PI', 'SST', 'RH', 'Vshear')
    for dsname in dsnames:
        for obs in ('ERA5', 'MERRA2'):
            get_obs_data(dsname, obs)

    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{daname}_season{len_season}mon.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()