#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Fri Dec 31 12:49:58 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
import xlinregress
import xfilter
#import xpyleoclim
import cartopy.crs as ccrs, cartopy.feature as cfeature
from misc.landmask import whereocean
from misc.seasons import sel_season
import geoxarray
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
FA = True if 'FA' in sys.argv else False
FAtrop = True if 'FAtrop' in sys.argv else False
Obs = True if 'Obs' in sys.argv else False
genesis = True if 'genesis' in sys.argv else False
doLowpass = True if 'lowpass' in sys.argv else False
if doLowpass:
    n_window = 40
    lowpass = lambda da: da.filter.lowpass(1/n_window, dim='year', padtype='even')
else:
    lowpass = lambda da: da
#params from sys.argv
#params for figure
show_slope = True if len(sys.argv)>1 and 'slope' in sys.argv[1:] else False #whether to plot slope of regression or corrcoef
years = slice(2001,2300) #2400)
if Obs: years = slice(1980,2023)
nyears = years.stop - years.start + 1
nino34Season = 'ASO'

"""
#daname = 'sst'
#n_window = (500, 40)
#bandpass = lambda da: da.filter.bandpass((1/n_window[0], 1/n_window[1]), dim='year', padtype='even')
n_window = 40
lowpass = lambda da: da.filter.lowpass(1/n_window, dim='year', padtype='even')
lowpass = lambda da: da
years = slice(1501,2000)
ifile_sst = '/tigress/wenchang/MODEL_OUT/FLOR/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/analysis/sst.0001-2000.yearly.nc'
ifile_tc = '/tigress/wenchang/analysis/TC/FLOR/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.TS17.0001-2000.yearly.nc'
"""

if FA:
    ifile_nino34 = 't_surf_FLOR_CTL1860_v202407_FA_tigercpu_intelmpi_18_576PE_1901-2400_nino34.nc'
    ifile_tc = 'tc_density.FLOR.CTL1860_v202407_FA_tigercpu_intelmpi_18_576PE.TS17.1901-2400.yearly.nc'
elif FAtrop:
    ifile_nino34 = 't_surf_FLOR_CTL1860_v202407_FAtrop_tigercpu_intelmpi_18_576PE_1901-2314_nino34.nc'
    ifile_tc = 'tc_density.FLOR.CTL1860_v202407_FAtrop_tigercpu_intelmpi_18_576PE.TS17.1901-2314.yearly.nc'
else:
    ifile_nino34 = 't_surf_FLOR_CTL1860_v202407_tigercpu_intelmpi_18_576PE_1901-2500_nino34.nc'
    ifile_tc = 'tc_density.FLOR.CTL1860_v202407_tigercpu_intelmpi_18_576PE.TS17.1901-2500.yearly.nc'
if Obs:
    ifile_nino34 = 'HadISST_sst.187001-202507.nc'
    ifile_tc = 'IBTrACS.ALL.v04r01.2024-10-24.density.TS17.1980-2023.yearly.nc'

ofile = f'linregress_map_FLOR_TC_nino34_{years.start}-{years.stop}.nc'

#if os.path.exists(ofile):
#    ds = xr.open_dataset(ofile)
#    print('[loaded]:', ofile)
#else:
if True:
    #TC
    print(ifile_tc)
    if genesis:
        da = xr.open_dataset(ifile_tc)['density_g'].pipe(lowpass).sel(year=years)
    else:
        da = xr.open_dataset(ifile_tc)['density'].pipe(lowpass).sel(year=years)
    da_tc = da#.pipe(whereocean)
    #da_tc = ( da - da.mean() )/da.std() #standardize
    #SST
    print(ifile_nino34)
    #da = xr.open_dataset(ifile_sst)['sst']
    if Obs:
        da = xr.open_dataset(ifile_nino34)['sst'].sel(longitude=slice(-170,120), latitude=slice(5,-5)).geo.fldmean() \
            .pipe(sel_season, nino34Season).groupby('time.year').mean('time')
    else:
        da = xr.open_dataset(ifile_nino34)['t_surf'].pipe(sel_season, nino34Season).groupby('time.year').mean('time')
    #da = ds_lmr[daname].groupby('time.year').mean('time').load()
    da_nino34 = da.pipe(lowpass).sel(year=years)
    da_nino34 = ( da_nino34 - da_nino34.mean() )/da_nino34.std() #standardize
    if True:
        print('linear regression...')
        """
        rg = da_sst.linregress.on(da_tc, dim='year', ess_on=True)
        ds = rg[['slope', 'r', 'dof', 'pvalue']]
        ds['slope'].attrs['units'] = 'K per STD HU'
        """
        rg = da_tc.linregress.on(da_nino34, dim='year', ess_on=True)
        ds = rg[['slope', 'r', 'dof', 'pvalue']]
        ds['slope'].attrs['units'] = 'TC days per year per 10x10deg box per STD Nino3.4 index'
    #ds.to_netcdf(ofile)
    #print('[saved]:', ofile)
 
def wyplot(ds, show_slope=False, ax=None, **kws):
    if ax is None: fig, ax = plt.subplots(subplot_kw=dict(projection=ccrs.Robinson(central_longitude=180)))
    dproj = ccrs.PlateCarree()
    land_color = 'k'
    if show_slope:
        da = ds.slope
        levels = np.arange(-0.12, 0.121, 0.02)
    else:
        da = ds.r
        levels = np.arange(-1, 1.01, 0.1)
    p = ds.pvalue
    da.plot.contourf(ax=ax, transform=dproj, levels=levels, **kws)
    #mapplot(ax=ax, fill_continents=True)
    da.where(p<0.05).plot.contourf(ax=ax, transform=dproj, colors='none', hatches=['...'], add_colorbar=False)

    #ax.add_feature(cfeature.LAND, color=land_color)
    ax.coastlines(linewidth=0.5, color='gray')
    ax.set_global()
    #ax.gridlines()
    lons = 360-170, 360-120
    lats = -5, 5
    ax.plot([lons[0], lons[1], lons[1], lons[0], lons[0]],
        [lats[0], lats[0], lats[1], lats[1], lats[0]], 
        transform=dproj, color='gray')
    
    if show_slope:
        title = f'TC track density due to {nino34Season} Nino34 ({nyears} years)'
    else:
        title = f'corr_coef(TC track density, {nino34Season} Nino34), {nyears} years'
    if Obs: 
        title += ', Obs'
    elif FAtrop: 
        title += ', FAtrop'
    elif FA: 
        title += ', FA'
    if doLowpass:
        title += f', {n_window}lp'
    if genesis: title = title.replace('track density', 'genesis')
    ax.set_title('')
    ax.set_title(title, loc='center') 

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    from geoplots import mapplot
    import cartopy.crs as ccrs, cartopy.feature as cfeature
    plt.close()
    proj = ccrs.Robinson(central_longitude=180)
    #proj = ccrs.Mollweide(central_longitude=180)
    dproj = ccrs.PlateCarree()
    figsize = (6,3)
    fig, ax = plt.subplots(figsize=figsize, subplot_kw=dict(projection=proj))
    wyplot(ds, show_slope=show_slope, ax=ax)

    """
    figsize = (6,6)
    fig, axes = plt.subplots(2, 1, figsize=figsize, subplot_kw=dict(projection=proj))

    ax = axes[0]
    wyplot(ds, show_slope=False, ax=ax)

    ax= axes[1]
    wyplot(ds, show_slope=True, ax=ax)
    """

    #savefig
    if 'savefig' in sys.argv or 's' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        if show_slope:
            figname = figname.replace('.png', '__slope.png')
        else:
            figname = figname.replace('.png', '__corr.png')
        if Obs: 
            figname = figname.replace('.png', '_Obs.png')
        elif FAtrop: 
            figname = figname.replace('.png', '_FAtrop.png')
        elif FA: 
            figname = figname.replace('.png', '_FA.png')
        if genesis: figname = figname.replace('.png', '_genesis.png')
        if doLowpass: figname = figname.replace('.png', f'_{n_window}lowpass.png')
        if 'overwritefig' in sys.argv or 'o' in sys.argv:
            wysavefig(figname, overwritefig=True)
        else:
            wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()