#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Oct  3 11:13:36 EDT 2019
from datetime import datetime
import os.path, sys, os
import xarray as xr, numpy as np, pandas as pd
import matplotlib.pyplot as plt
from misc.cim import cim

plt.close()
basin = 'NA'
tag = 'sstModel'
figname = f'fig.tcCounts.anom.{tag}.{basin}.png'
figsize = (8, 4)#None
fig = plt.figure(figsize=figsize)
years = range(1,6) 
xticks = years
ylabel = f'%'
xlabel = 'year'
title = f'Response of {basin} TCs to Large Volcanic Eruptions: SST model'
alpha = 0.3
capsize = 3 
fmt = 'o'
bar_width = 0.1
ecolor = 'gray'

volc = 'CTL'
color = 'k'
ifile = '/tigress/wenchang/analysis/TC_sst_model/CTL1860_noleap_tigercpu_intelmpi_18_576PE.TCfit.0011-0044.nc'
with xr.open_dataset(ifile) as ds:
    sstMDR = ds['sstMDR']
    sstTrop = ds['sstTrop']
# transform to ensembles
sstMDRctl = []
sstTropCtl = []
for i in range(30):
    sstMDRctl.append(sstMDR.isel(year=slice(i, i+5)).assign_coords(year=years))
    sstTropCtl.append(sstTrop.isel(year=slice(i, i+5)).assign_coords(year=years))
sstMDRctl = xr.concat(sstMDRctl, dim=pd.Index(range(1, 31), name='en') )
sstTropCtl = xr.concat(sstTropCtl, dim=pd.Index(range(1, 31), name='en') )

volc = 'Pinatubo'
year_s = 1991
color = 'C2'
offset = 0.1*0
ifile = f'/tigress/wenchang/analysis/TC_sst_model/{volc}_PI_ens_noleap.TCfit.ens01-30.{year_s}-{year_s+4}.nc'
sstMDR = xr.open_dataset(ifile)['sstMDR'].assign_coords(year=years)
sstTrop = xr.open_dataset(ifile)['sstTrop'].assign_coords(year=years)
da = ( 1.388*(sstMDR - sstMDRctl) - 1.521*(sstTrop - sstTropCtl) )
err = da.std(dim='en')/da.en.size**0.5 * 100
anom = da.mean('en') * 100 # percent change
plt.bar(da.year+offset, anom, yerr=err, color=color, capsize=capsize, label=volc, width=bar_width, ecolor=ecolor)

volc = 'Agung'
year_s = 1963
color = 'C0'
offset = 0.1
ifile = f'/tigress/wenchang/analysis/TC_sst_model/{volc}_PI_ens_noleap.TCfit.ens01-30.{year_s}-{year_s+4}.nc'
sstMDR = xr.open_dataset(ifile)['sstMDR'].assign_coords(year=years)
sstTrop = xr.open_dataset(ifile)['sstTrop'].assign_coords(year=years)
da = ( 1.388*(sstMDR - sstMDRctl) - 1.521*(sstTrop - sstTropCtl) )
err = da.std(dim='en')/da.en.size**0.5 * 100
anom = da.mean('en') * 100 # percent change
plt.bar(da.year+offset, anom, yerr=err, color=color, capsize=capsize, label=volc, width=bar_width, ecolor=ecolor)

volc = 'Chichon'
year_s = 1982
color = 'C1'
offset = 0.1*2
ifile = f'/tigress/wenchang/analysis/TC_sst_model/{volc}_PI_ens_noleap.TCfit.ens01-30.{year_s}-{year_s+4}.nc'
sstMDR = xr.open_dataset(ifile)['sstMDR'].assign_coords(year=years)
sstTrop = xr.open_dataset(ifile)['sstTrop'].assign_coords(year=years)
da = ( 1.388*(sstMDR - sstMDRctl) - 1.521*(sstTrop - sstTropCtl) )
err = da.std(dim='en')/da.en.size**0.5 * 100
anom = da.mean('en') * 100 # percent change
plt.bar(da.year+offset, anom, yerr=err, color=color, capsize=capsize, label=volc, width=bar_width, ecolor=ecolor)

volc = 'Chichon_CMIP5volc'
year_s = 1982
color = 'C1'
offset = 0.1*3
ifile = f'/tigress/wenchang/analysis/TC_sst_model/{volc}_PI_ens_noleap.TCfit.ens01-30.{year_s}-{year_s+4}.nc'
sstMDR = xr.open_dataset(ifile)['sstMDR'].assign_coords(year=years)
sstTrop = xr.open_dataset(ifile)['sstTrop'].assign_coords(year=years)
da = ( 1.388*(sstMDR - sstMDRctl) - 1.521*(sstTrop - sstTropCtl) )
err = da.std(dim='en')/da.en.size**0.5 * 100
anom = da.mean('en') * 100 # percent change
plt.bar(da.year+offset, anom, yerr=err, color=color, capsize=capsize, label=volc, width=bar_width, ecolor=ecolor, hatch='///')

volc = 'StMaria'
year_s = 1902
color = 'C3'
offset = 0.1*4
ifile = f'/tigress/wenchang/analysis/TC_sst_model/{volc}_PI_ens_noleap.TCfit.ens01-30.{year_s}-{year_s+4}.nc'
sstMDR = xr.open_dataset(ifile)['sstMDR'].assign_coords(year=years)
sstTrop = xr.open_dataset(ifile)['sstTrop'].assign_coords(year=years)
da = ( 1.388*(sstMDR - sstMDRctl) - 1.521*(sstTrop - sstTropCtl) )
err = da.std(dim='en')/da.en.size**0.5 * 100
anom = da.mean('en') * 100 # percent change
plt.bar(da.year+offset, anom, yerr=err, color=color, capsize=capsize, label=volc, width=bar_width, ecolor=ecolor)

plt.axhline(0, color='gray', ls='--')
if xticks is not None:
    plt.xticks(xticks)
if ylabel is not None:
    plt.ylabel(ylabel)
if xlabel is not None:
    plt.xlabel(xlabel)
if title is not None:
    plt.title(title)
#plt.legend(ncol=4, bbox_to_anchor=(0, 1), loc='lower left')
plt.legend(ncol=3, loc='lower right')

if figname is not None:
    plt.tight_layout()
    plt.savefig(figname)
plt.show()