#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Jun  9 13:58:05 EDT 2021
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')
import sys, os.path, os, glob, datetime
import xarray as xr, numpy as np, pandas as pd, matplotlib.pyplot as plt
#more imports
import geoxarray
import xfilter
from misc.cim import cim
from geoplots import mapplot
#
if __name__ == '__main__':
    tt.check('end import')
#
#start from here
datafile = __file__.replace('.py', '.nc')
if os.path.exists(datafile):
    ds = xr.open_dataset(datafile)
    print('[loaded]:', datafile)
else:
    volcs = ['Pinatubo', 'Agung', 'StMaria']
    #Ts
    #ctl
    ifile_ctl = '/tigress/wenchang/analysis/TC/CTL1860_noleap_tigercpu_intelmpi_18_576PE/netcdf/tc_density.TS17.0011-0044.nc'
    ctl = xr.open_dataset(ifile_ctl)['density'].drop('time')
    units = ctl.attrs['units']
    tcmask = ctl.mean('time') > 0 #0-1 tc mask
    #convert to ctl ensemble
    ctl_ens = []
    for ii in range(1,30*12+1, 12):
        da_ = ctl.isel(time=slice(ii-1, ii-1+5*12)) 
        ctl_ens.append(da_)
    ctl = xr.concat(ctl_ens, dim=pd.Index(range(1,31), name='en'))

    #volc
    das_volcs = []
    for volc in volcs:
        print(volc)
        if volc == 'Pinatubo':
            erupt_month = 6
            erupt_year = 1991
        elif volc == 'Agung':
            erupt_month = 3
            erupt_year = 1963
        elif volc == 'StMaria':
            erupt_month = 10
            erupt_year = 1902
        ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_density.TS17.nens30.{erupt_year}-{erupt_year+4}.nc'
        da = xr.open_dataset(ifile)['density'] #volc run
        daa = da - ctl #anomaly from the control run
        da_ = daa.isel(time=slice(erupt_month-1, erupt_month-1+3*12)).mean('time')*12 #average over 3 yearas after eruption; units from per month to per year
        das_volcs.append(da_)
    #concat along volcs
    da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units='TC days per year per 10x10deg box')
    tcdensity_anom = da
     
    ds = xr.Dataset(dict(tcdensity_anom=tcdensity_anom, tcmask=tcmask))
    ds.to_netcdf(datafile)
    print(datafile)

def axscale(ax=None, right=0, left=0, up=0, down=0):
    '''scale the axes by the four sides'''
    if ax is None:
        ax = plt.gca()
    x0, y0, w, h = ax.get_position().bounds
    w_new = w*(1 + right + left)
    h_new = h*(1 + up + down)
    x0_new = x0 - w*left
    y0_new = y0 - h*down
    ax.set_position([x0_new, y0_new, w_new, h_new])
    return ax

def wyplot(ax=None, volc='Pinatubo', lowpass=None, map_on=True, colorbar_label_on=True):
    vname = 'tcdensity_anom'
    lv = 2**np.arange(-3, 4.1)
    levels = list(-lv[-1::-1]) + [0,] + list(lv)
    coastline_width = .5
    cbar_ticks = (-16, -4, -1, -0.25, 0, 0.25, 1, 4, 16)
    cbar_format = '%.2g'
    alpha=0.2
    hatches = ['///']

    if ax is None:
        fig, ax = plt.subplots()
    if lowpass is None:
        lowpass = lambda x: x.filter.lowpass(1/5, dim='time', padtype='constant')
    alpha = 0.2
    if volc == 'Pinatubo':
        color = 'k'
    elif volc == 'Agung':
        color = 'C0'
    elif volc == 'StMaria':
        color = 'C1'
    da = ds[vname].sel(volc=volc)
    units = 'TC days per year per 10x10deg'
    #long_name = 'TC track density'
    if map_on:
        damean = da.mean('en')
        uncertainty = da.pipe(cim, dim='en')
        sigmask = np.abs(damean) > uncertainty
        tcmask = ds.tcmask
        if colorbar_label_on:
            damean.where(tcmask).rename(units).plot(ax=ax, levels=levels, cbar_kwargs={'ticks': cbar_ticks, 'format': cbar_format}, rasterized=True)
        else:
            damean.where(tcmask).rename('').plot(ax=ax, levels=levels, cbar_kwargs={'ticks': cbar_ticks, 'format': cbar_format}, rasterized=True)
        damean.where(tcmask).where(sigmask).pipe(lambda x: x*0).plot.contourf(ax=ax, colors='none', hatches=hatches, add_colorbar=False)
        mapplot(ax=ax, lon=da.lon, lat=(-50, 50), linewidth=coastline_width)

        ax.set_xlabel('')
        ax.set_ylabel('')
    else:
        da = da.mean('lon') 
        damean = da.mean('en')
        uncertainty = da.pipe(cim, dim='en')
        ci = damean - uncertainty, damean + uncertainty
        ax.fill_betweenx(da.lat, ci[0], ci[-1], color=color, alpha=alpha)
        damean.plot(ax=ax, y='lat', lw=1.5, color=color)

        ax.axhline(0, color='gray', lw=1, ls='--')
        ax.axvline(0, color='gray', lw=1, ls='--')
        ax.set_xlabel('Zonal mean track density response')
        ax.set_ylabel('')

    ax.set_yticks(np.arange(-45,46,15))
    ax.set_yticklabels(['45$^\circ$S', '30$^\circ$S', '15$^\circ$S', '0', '15$^\circ$N', '30$^\circ$N', '45$^\circ$N'])

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    constrained_layout_off()
    fig, axes = plt.subplots(3, 2, figsize=(8, 8), sharey=True, sharex='col', dpi=100)
    ax = axes[0,0]
    wyplot(ax=ax, volc='Pinatubo', map_on=False)
    ax = axes[0,1]
    wyplot(ax=ax, volc='Pinatubo', map_on=True, colorbar_label_on=False)
    ax = axes[1,0]
    wyplot(ax=ax, volc='Agung', map_on=False)
    ax = axes[1,1]
    wyplot(ax=ax, volc='Agung', map_on=True)
    ax = axes[2,0]
    wyplot(ax=ax, volc='StMaria', map_on=False)
    ax = axes[2,1]
    wyplot(ax=ax, volc='StMaria', map_on=True, colorbar_label_on=False)

    for ax,label in zip(axes.flat, list('ABCDEF')):
        ax.set_title('')
        ax.text(0.01, 0.99, label, transform=ax.transAxes, ha='left', va='top', fontweight='bold')
    for ax in axes[0:2, 0]:
        ax.set_xlabel('')
    for ax,volc in zip(axes[:, 0], ds.volc.values):
        ax.set_title(volc, loc='left')

    plt.tight_layout()
    for ax in axes[:, 0]:
        axscale(ax, right=-0.5)
    for ax in axes[:, 1]:
        axscale(ax, left=0.6)

    #savefig
    if 'savefig' in sys.argv:
        figname = __file__.replace('.py', f'.png')
        wysavefig(figname)
    tt.check(f'**Done**')
    plt.show()
    constrained_layout_on()