#!/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):
    #dsname = 'PI'
    print(dsname)
    if dsname == 'PI_vmax':
        dsname_new = 'PI'
        long_name = 'Potential Intensity'
        units = 'm/s'
    elif dsname == 'skin_temperature':
        dsname_new = 'SST'
        long_name = 'SST'
        units = 'K'
    elif dsname == 'relative_humidity':
        dsname_new = 'RH'
        long_name = 'Relative Humidity'
        units = '%'
    elif dsname == 'vertical_wind_shear':
        dsname_new = 'Vshear'
        long_name = 'Vertical Wind Shear'
        units = 'm/s'
    daname = dsname
    trend_units = f'{units}/year'
    ifile_NH = glob.glob(f'ERA5MERRA2/{obs}*{dsname}*_NH_*.nc')[0]
    ifile_SH = glob.glob(f'ERA5MERRA2/{obs}*{dsname}*_SH_*.nc')[0]
    len_season = 3
    print(f'{ifile_NH = }')
    print(f'{ifile_SH = }')
    #
    da = xr.open_dataset(ifile_NH)[daname].groupby('time.year').mean('time').drop(['lon', 'lat'])*wgts[0] \
        + xr.open_dataset(ifile_SH)[daname].groupby('time.year').mean('time').drop(['lon', 'lat'])*wgts[1]
    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}.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_vmax', 'skin_temperature', 'relative_humidity', 'vertical_wind_shear')
    obss = ('ERA5', 'MERRA2')
    for dsname in dsnames:
        for obs in obss:
            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()