#!/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
#
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']
    #F
    daname = 'netrad_toa'
    print(daname)
    ifile_ctl = f'data/nudgeclimo_all_model_CTL1860_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc'
    ctl = xr.open_dataarray(ifile_ctl)
    das_volcs = []
    for volc in volcs:
        ifile = f'data/{volc}_ens_noleap_nudgeclimo_all_model1860.atmos_month.{daname}.nc'
        da = xr.open_dataarray(ifile) #volc run
        units = da.attrs['units']
        #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run
        daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run
        daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS'))
        #global mean
        gl = daa.geo.fldmean()
        #NH mean
        nh = daa.sel(lat=slice(0,90)).geo.fldmean()
        #SH mean
        sh = daa.sel(lat=slice(-90,0)).geo.fldmean()
        #concat along gl/nh/sh
        da_ = xr.concat([gl,nh,sh], dim=pd.Index(['GL', 'NH', 'SH'], name='sphere'))
        das_volcs.append(da_)
    #concat along volcs
    da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units=units)
    F = da
    
    #N
    daname = 'netrad_toa'
    print(daname)
    ifile_ctl = f'data/CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc'
    ctl = xr.open_dataarray(ifile_ctl)
    das_volcs = []
    for volc in volcs:
        ifile = f'data/{volc}_PI_ens_noleap.atmos_month.{daname}.nc'
        da = xr.open_dataarray(ifile) #volc run
        units = da.attrs['units']
        #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run
        daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run
        daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS'))
        #global mean
        gl = daa.geo.fldmean()
        #NH mean
        nh = daa.sel(lat=slice(0,90)).geo.fldmean()
        #SH mean
        sh = daa.sel(lat=slice(-90,0)).geo.fldmean()
        #concat along gl/nh/sh
        da_ = xr.concat([gl,nh,sh], dim=pd.Index(['GL', 'NH', 'SH'], name='sphere'))
        das_volcs.append(da_)
    #concat along volcs
    da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units=units)
    N = da

    #Ts
    daname = 't_surf'
    print(daname)
    ifile_ctl = f'data/CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc'
    ctl = xr.open_dataarray(ifile_ctl)
    das_volcs = []
    for volc in volcs:
        ifile = f'data/{volc}_PI_ens_noleap.atmos_month.{daname}.nc'
        da = xr.open_dataarray(ifile) #volc run
        units = da.attrs['units']
        #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run
        daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run
        daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS'))
        #global mean
        gl = daa.geo.fldmean()
        #NH mean
        nh = daa.sel(lat=slice(0,90)).geo.fldmean()
        #SH mean
        sh = daa.sel(lat=slice(-90,0)).geo.fldmean()
        #concat along gl/nh/sh
        da_ = xr.concat([gl,nh,sh], dim=pd.Index(['GL', 'NH', 'SH'], name='sphere'))
        das_volcs.append(da_)
    #concat along volcs
    da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).assign_attrs(units=units)
    tsa = da
    
    #precip
    daname = 'precip'
    print(daname)
    ifile_ctl = f'data/CTL1860_noleap_tigercpu_intelmpi_18_576PE.atmos_month.{daname}.nc'
    ctl = xr.open_dataarray(ifile_ctl)
    das_volcs = []
    for volc in volcs:
        ifile = f'data/{volc}_PI_ens_noleap.atmos_month.{daname}.nc'
        da = xr.open_dataarray(ifile) #volc run
        #units = da.attrs['units']
        #daa = da - ctl.assign_coords(time=da.time.values) #anomaly from the control run
        daa = da.assign_coords(time=ctl.time.values) - ctl #anomaly from the control run
        daa = daa.assign_coords(time=xr.cftime_range('0001-01', '0005-12', freq='MS'))
        #global mean
        gl = daa.geo.fldmean()
        #NH mean
        nh = daa.sel(lat=slice(0,90)).geo.fldmean()
        #SH mean
        sh = daa.sel(lat=slice(-90,0)).geo.fldmean()
        #concat along gl/nh/sh
        da_ = xr.concat([gl,nh,sh], dim=pd.Index(['GL', 'NH', 'SH'], name='sphere'))
        das_volcs.append(da_)
    #concat along volcs
    da = xr.concat(das_volcs, dim=pd.Index(volcs, name='volc')).pipe(lambda x: x*24*3600).assign_attrs(units='mm/day')#units from kg/m^2/s -> mm/day
    pra = da
     
    ds = xr.Dataset(dict(F=F, N=N, tsa=tsa, pra=pra))
    ds.to_netcdf(datafile)
    print(datafile)

def wyplot(ax=None, vname='tsa', volc='Pinatubo', lowpass=None, legend_on=True):
    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':
        time_volc = pd.date_range('1991-01', periods=60, freq='MS')
        date_erupt = time_volc[5]
    elif volc == 'Agung':
        time_volc = pd.date_range('1963-01', periods=60, freq='MS')
        date_erupt = time_volc[2]
    elif volc == 'StMaria':
        time_volc = pd.date_range('1902-01', periods=60, freq='MS')
        date_erupt = time_volc[9]
    if vname == 'F':
        units = 'W m$^{-2}$'
        long_name = 'F'
    elif vname == 'tsa':
        units = 'K'
        long_name = 'T$_s$'
    elif vname == 'pra':
        units = 'mm day$^{-1}$'
        long_name = 'T$_s$'
    da = ds[vname].sel(volc=volc).pipe(lowpass).assign_coords(time=time_volc)
    if vname == 'F':
        da = da.sel(en=range(1,11))
    damean = da.mean('en')
    dd = da.pipe(cim, dim='en')
    ci = damean - dd, damean + dd
    #global mean
    sphere, color = 'GL', 'k'
    ax.fill_between(da.time.values, ci[0].sel(sphere=sphere), ci[1].sel(sphere=sphere), color=color, alpha=alpha)
    damean.sel(sphere=sphere).plot(ax=ax, color=color, label='Global')
    #SH mean
    sphere, color = 'SH', 'C0'
    ax.fill_between(da.time.values, ci[0].sel(sphere=sphere), ci[1].sel(sphere=sphere), color=color, alpha=alpha)
    damean.sel(sphere=sphere).plot(ax=ax, color=color, label='SH')
    #NH mean
    sphere, color = 'NH', 'C3'
    ax.fill_between(da.time.values, ci[0].sel(sphere=sphere), ci[1].sel(sphere=sphere), color=color, alpha=alpha)
    damean.sel(sphere=sphere).assign_attrs(units=units, long_name=long_name).plot(ax=ax, color=color, label='NH')

    ax.axhline(0, color='gray', ls='--')
    ax.axvline(date_erupt, color='gray', ls='--')
    if legend_on:
        ax.legend(frameon=False)
    ax.set_title(volc)
    ax.set_xlabel('')
    #ax.set_xlim(time_volc[0], time_volc[-1])

if __name__ == '__main__':
    from wyconfig import * #my plot settings
    fig,axes = plt.subplots(3, 3, figsize=(10,6), sharex='col', sharey='row')
    for ii,vname in enumerate(['F', 'tsa', 'pra']):
        for jj,volc in enumerate(ds.volc.values):
            ax = axes[ii,jj]
            if ii==0 and jj==2:
                legend_on = True
            else:
                legend_on = False
            wyplot(ax=ax, vname=vname, volc=volc, legend_on=legend_on)
            if ii > 0:
                ax.set_title('')
    for ax,label in zip(axes.flat, list('abcdefghi'.upper())):
        ax.text(1-0.01,1-0.02,label, ha='right', va='top', fontweight='bold', transform=ax.transAxes)

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