#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Aug 25 12:03:14 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob
#import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
from geoplots import mapplot
from mystats import p2t
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
season = ['MAM', 'DJF'][0]
alpha = 0.1 # significance level
n_sample = 10 # top/bottom n_sample NAO index years 
year_range = range(1980,2020)
ifile = 'norm.nao.monthly.b5001.current.ascii.table'

df = pd.read_csv(ifile, sep=r'\s+', header=None, index_col=0, names=['year',]+ list(range(1,13)))
ts = xr.DataArray(df.values.ravel(), dims='time', coords=[pd.date_range('1950-01', end='2020-12', freq='MS')])
if season == 'DJF':
    ts_season = ts.where((month==12)|(month==1)|(month==2)) \
        .resample(time='A-Feb').mean()
elif season == 'MAM':
    ts_season = ts.where((month==3)|(month==4)|(month==5)) \
        .resample(time='A-May').mean()
ts_season = ts_season.groupby('time.year').mean('time') \
    .sel(year=year_range)
years_naoP = ts_season.sortby(ts_season).year[-n_sample:]
years_naoN = ts_season.sortby(ts_season).year[:n_sample]

#ERA5 q2m
ifiles = '/home/wenchang/era5/analysis/2m_specific_humidity/monthly/*.nc'
ds = xr.open_mfdataset(ifiles)
ds.load()
month = ds.time.dt.month
da = ds.q2m
if season == 'DJF':
    da = da.isel(time=(month==12)|(month==1)|(month==2))
    da = da.resample(time='A-Feb').mean().groupby('time.year').mean('time').sel(year=range(1980, 2020))
elif season == 'MAM':
    da = da.isel(time=(month==3)|(month==3)|(month==5))
    da = da.resample(time='A-May').mean().groupby('time.year').mean('time').sel(year=range(1980, 2020))
daa = da - da.mean('year')
spread = p2t(alpha, df=n_sample-1) * daa.std('year') * n_sample**(-1/2)

# recenter at lon=0
if daa.longitude.max() > 180:
    daa = daa.roll(longitude=daa.longitude.size//2).pipe(lambda x: x.assign_coords(longitude=x.longitude.where(x.longitude<180, other=x.longitude-360)))
if spread.longitude.max() > 180:
    spread = spread.roll(longitude=daa.longitude.size//2).pipe(lambda x: x.assign_coords(longitude=x.longitude.where(x.longitude<180, other=x.longitude-360)))

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    figname = __file__.replace('.py', f'_{season}_{tt.today()}.png')
    close('all')
    figsize = (6, 8)
    fig, axes = plt.subplots(3, 1, figsize=figsize)
    ax = axes[0]
    plt.sca(ax)
    ts_season.plot(color='k', ax=ax)
    plt.plot(years_naoN, ts_season.sel(year=years_naoN), '*', color='C0', label=f'{n_sample} lowest NAO years')
    plt.plot(years_naoP, ts_season.sel(year=years_naoP), '*', color='C1', label=f'{n_sample} highest NAO years')
    ax.set_title(f'{season} NAO index', loc='left')
    ax.set_xlabel('')
    ax.axhline(0, ls='--', color='gray')
    
    ax = axes[1]
    plt.sca(ax)
    daa.sel(year=years_naoP).mean('year').pipe(lambda x: x*1000).assign_attrs(units='g/kg').plot(cmap='BrBG', levels=21)
    daa.sel(year=years_naoP).mean('year').pipe(lambda x: x.where(np.abs(x)>spread)*0).plot.contourf(colors='none', hatches=['///'], add_colorbar=False)
    mapplot()
    #ax.set_title('NAO+', loc='left')
    ax.text(-180, 90, 'NAO+', ha='left', va='bottom', color='C1')
    ax.set_xlabel(''); ax.set_ylabel('')

    ax = axes[2]
    plt.sca(ax)
    daa.sel(year=years_naoN).mean('year').pipe(lambda x: x*1000).assign_attrs(units='g/kg').plot(cmap='BrBG', levels=21)
    daa.sel(year=years_naoN).mean('year').pipe(lambda x: x.where(np.abs(x)>spread)*0).plot.contourf(colors='none', hatches=['///'], add_colorbar=False)
    mapplot()
    #ax.set_title('NAO-', loc='left')
    ax.text(-180, 90, 'NAO-', ha='left', va='bottom', color='C0')
    ax.set_xlabel(''); ax.set_ylabel('')

    plt.savefig(figname)
    print('[saved]:', figname)
    tt.check(f'**Done**')
    plt.show()