#!/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
daname = 'Vshear'
long_name = 'Vertical Wind Shear'
units = 'm/s'
trend_units = 'm/s/year'
len_season = 3
if len_season == 3:
    ifile_nat_NH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-nat_NH_ASO_') and f.endswith('.nc')][0]
    ifile_nat_SH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-nat_SH_FMA_') and f.endswith('.nc')][0]
    ifile_ghg_NH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-GHG_NH_ASO_') and f.endswith('.nc')][0]
    ifile_ghg_SH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-GHG_SH_FMA_') and f.endswith('.nc')][0]
elif len_season == 5:
    ifile_nat_NH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-nat_NH_JJASO_') and f.endswith('.nc')][0]
    ifile_nat_SH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-nat_SH_DJFMA_') and f.endswith('.nc')][0]
    ifile_ghg_NH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-GHG_NH_JJASO_') and f.endswith('.nc')][0]
    ifile_ghg_SH = [f for f in os.listdir() if f.startswith(f'{daname}_hist-GHG_SH_DJFMA_') and f.endswith('.nc')][0]
print(f'{ifile_nat_NH = }')
print(f'{ifile_nat_SH = }')
print(f'{ifile_ghg_NH = }')
print(f'{ifile_ghg_SH = }')
#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)
#nat
da_nat = xr.open_dataarray(ifile_nat_NH)*wgts[0] + xr.open_dataarray(ifile_nat_SH)*wgts[1]
da_nat = da_nat.sel(year=slice(1982, 2014))
lintrend_nat = da_nat.linregress.on(da_nat.year)
#ghg
da_ghg = xr.open_dataarray(ifile_ghg_NH)*wgts[0] + xr.open_dataarray(ifile_ghg_SH)*wgts[1]
da_ghg = da_ghg.sel(year=slice(1982, 2014))
lintrend_ghg = da_ghg.linregress.on(da_ghg.year)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from scipy import stats
    plt.close('all')
    alpha = 0.5
    bins = 10
    density = True
    fig,axes = plt.subplots(2, 1, figsize=(6,6))
    ax = axes[0]
    #pdfs
    lintrend_nat.slope.plot.hist(ax=ax, density=density, alpha=alpha)
    lintrend_ghg.slope.plot.hist(ax=ax, density=density, alpha=alpha)

    #norm fit
    mu_nat, sigma_nat = stats.norm.fit(lintrend_nat.slope)
    mu_ghg, sigma_ghg = stats.norm.fit(lintrend_ghg.slope)
    xmin = min(lintrend_nat.slope.min(), lintrend_ghg.slope.min()).item()
    xmax = max(lintrend_nat.slope.max(), lintrend_ghg.slope.max()).item()
    xran = xmax - xmin
    x = np.linspace(xmin-xran*0.1, xmax+xran*0.1, 100)
    ax.plot(x, stats.norm.pdf(x, mu_nat, sigma_nat), color='C0', label='hist-nat')
    ax.plot(x, stats.norm.pdf(x, mu_ghg, sigma_ghg), color='C1', label='hist-GHG')

    #ttest_int, kstest
    ttest = stats.ttest_ind(lintrend_nat.slope, lintrend_ghg.slope)
    kstest = stats.kstest(lintrend_nat.slope, lintrend_ghg.slope)

    #
    ax.legend()
    ax.set_title(f'{long_name}, t-test p={ttest.pvalue:.3g}, KS-test p={kstest.pvalue:.3g}')
    ax.set_xlabel(trend_units)
    ax.set_ylabel('Probability Density')

    #time series
    ax = axes[1]
    da = da_nat.groupby(da_nat.model).mean('model_member').mean('model')
    da.plot(ax=ax, color='C0', marker='o', fillstyle='none', label='hist-nat')
    reg = da.linregress.on(da.year)
    reg.predicted.plot(ax=ax, color='C0', ls='--', label=f'hist-nat trend ({reg.slope.item():.2g} {trend_units})')
    da = da_ghg.groupby(da_ghg.model).mean('model_member').mean('model')
    da.plot(ax=ax, color='C1', marker='o', fillstyle='none', label='hist-GHG')
    reg = da.linregress.on(da.year)
    reg.predicted.plot(ax=ax, color='C1', ls='--', label=f'hist-GHG trend ({reg.slope.item():.2g} {trend_units})')

    ax.legend()
    ax.set_ylabel(units)

    
    #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()