#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Mon Mar 28 19:43:11 EDT 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
from scipy.stats import ttest_ind
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
season = ['JJA', 'SON', 'MAM', 'DJF', 'annual'][-2]# 'JJA'
if season == 'JJA':
    sel_season = lambda da: da.isel(time=(da.time.dt.month>=6)&(da.time.dt.month<=8))
elif season == 'DJF':
    def sel_season(da):
        time_new = da.indexes['time'].shift(1, 'MS')
        da['time'] = time_new
        da = da.isel(time=(da.time.dt.month>=1)&(da.time.dt.month<=3))
        return da
elif season == 'MAM':
    sel_season = lambda da: da.isel(time=(da.time.dt.month>=3)&(da.time.dt.month<=5))
elif season == 'SON':
    sel_season = lambda da: da.isel(time=(da.time.dt.month>=9)&(da.time.dt.month<=11))
else:
    sel_season = lambda da: da
years = slice('1990', '2014')
if len(sys.argv)>1:
    if '1990-2014' in sys.argv:
        years = slice('1990', '2014')
    elif '1999-2014' in sys.argv:
        years = slice('1999', '2014')
    elif '1976-1998' in sys.argv:
        years = slice('1976', '1998')
    elif '1901-2020' in sys.argv:
        years = slice('1901', '2020')
    elif '1990-2016' in sys.argv:
        years = slice('1990', '2016')
    else:
        years = slice('1990', '2014')
region = 'NA'
if len(sys.argv)>1:
    if 'global' in sys.argv:
        region = 'global'
if region == 'NA':
    sel_region = lambda da: da.sel(lon=slice(360-170, 360-50), lat=slice(10, 80))
elif region == 'global':
    sel_region = lambda da: da
case = 'urbanGlobal'
if len(sys.argv)>1:
    if 'urbanNA' in sys.argv:
        case = 'urbanNA'
    elif 'urbanAS' in sys.argv:
        case = 'urbanAS'
func_mean = lambda da: da.mean('en').mean('year')
new_ens = True
ten_ens = False#True
label = f'{case} - noUrban, {season}, {years.start}-{years.stop}, {region}'
if case == 'urbanGlobal' and new_ens:
    label += ', new ens'
    if ten_ens: label = label.replace('new ens', '10 ens')
tag = label.replace(',', '_').replace(' ', '')
pnwbox = 360-123, 360-119, 45, 52

ofile = __file__.replace('.py', f'_{tag}.nc')
ofile_pvalue = ofile.replace('.nc', '_pvalue.nc')
if False and os.path.exists(ofile_pvalue):
    daa = xr.open_dataarray(ofile)
    print('[loaded]:', ofile)
    pvalue = xr.open_dataarray(ofile_pvalue)
    print('[loaded]:', ofile_pvalue)
else:
    ifile = '../data_Tas_AM4urban_wasteCool_0urban_amip_5ens_1870-2020.nc'
    da0 = xr.open_dataarray(ifile).pipe(sel_season).sel(time=years).pipe(sel_region) \
        .load() \
        .groupby('time.year').mean('time')
    #more ens
    if ten_ens:
        ifile = './data_Tas_AM4urban_wasteCool_0urban_amip_new5ens_1990-2014.nc' 
        da0_more = xr.open_dataarray(ifile).pipe(sel_season).sel(time=years).pipe(sel_region) \
            .load() \
            .groupby('time.year').mean('time')
        da0 = xr.concat([da0, da0_more], dim='en')


    if case == 'urbanGlobal':
        ifile = '../data_Tas_AM4urban_wasteCool_amip_5ens_1870-2020.nc'
        da = xr.open_dataarray(ifile).pipe(sel_season).sel(time=years).pipe(sel_region) \
            .load() \
            .groupby('time.year').mean('time')
        if new_ens:
            da = da.sel(en=[0, 21, 22, 23, 24])
        else:
            da = da.sel(ens=[0, 11, 12, 13, 14])
        #more ens
        if ten_ens:
            ifile = './data_Tas_AM4urban_wasteCool_amip_new5ens_1990-2014.nc'
            da_more = xr.open_dataarray(ifile).pipe(sel_season).sel(time=years).pipe(sel_region) \
                .load() \
                .groupby('time.year').mean('time')
            da = xr.concat([da, da_more], dim='en')



    elif case == 'urbanNA':
        ifile = '../data_Tas_AM4urban_wasteCool_urbanNAonly_amip_5ens_1870-2020.nc'
        da = xr.open_dataarray(ifile).pipe(sel_season).sel(time=years).pipe(sel_region) \
            .load() \
            .groupby('time.year').mean('time')

    elif case == 'urbanAS':
        ifile = '../data_Tas_AM4urban_wasteCool_urbanASonly_amip_5ens_1870-2020.nc'
        da = xr.open_dataarray(ifile).pipe(sel_season).sel(time=years).pipe(sel_region) \
            .load() \
            .groupby('time.year').mean('time')

    daa = da.pipe(func_mean) - da0.pipe(func_mean)
    daa.to_dataset(name=da0.name).to_netcdf(ofile)
    print('[saved]:', ofile)

    #pvalue
    tvalue, pvalue = ttest_ind(da.stack(s=['en', 'year']), da0.stack(s=['en', 'year']), axis=-1)
    pvalue = xr.DataArray(pvalue, dims=da.stack(s=['en', 'year']).dims[:-1])
    for dim in pvalue.dims:
        pvalue = pvalue.assign_coords({dim: da[dim].values})
    pvalue.to_dataset(name='pvalue').to_netcdf(ofile_pvalue)
    print('[saved]:', ofile_pvalue)
 
 
if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    daa.plot(vmax=0.4, extend='both')
    mapplot()
    daa.where(pvalue<0.1).plot.contourf(colors='none', hatches=['......'], add_colorbar=False)

    lon0, lon1, lat0, lat1 = pnwbox
    plt.plot([lon0, lon1, lon1, lon0, lon0], [lat0, lat0, lat1, lat1, lat0], color='gray')

    plt.title(label)
    
    #savefig
    if len(sys.argv)>1 and 'savefig' in sys.argv[1:]:
        figname = __file__.replace('.py', f'_{tag}.png')
        if 'overwritefig' in sys.argv[1:]:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    print()
    plt.show()