#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Oct  3 11:13:36 EDT 2019
if __name__ == '__main__':
    from misc.timer import Timer
    tt = Timer(f'start {__file__}')

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 = 'global'
basin = 'NA'
#figname = f'fig.tcCounts.anom.{basin}.png'
figname = __file__.replace('.py', f'_{basin}_{tt.today()}.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'
alpha = 0.3
capsize = 3 
fmt = 'o'
bar_width = 0.1
ecolor = 'gray'

volc = 'CTL'
color = 'k'
#ifile = '/tigress/wenchang/analysis/TC/CTL1860_noleap_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0011-0044.yearly.nc'
ifile = '/tigress/wenchang/analysis/TC/CTL1860_noleap_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.TS17.0001-0100.yearly.nc'
da = xr.open_dataset(ifile)[basin]
da = da.sel(year=range(11,45))
# transform to ensembles
das = []
for i in range(30):
    das.append(da.isel(year=slice(i, i+5)).assign_coords(year=years))
da = xr.concat(das, dim=pd.Index(range(1, 31), name='en') )
meanCtl = da.mean('en')
ctl = da

volc = 'Pinatubo'
year_s = 1991
color = 'C2'
offset = 0.1*0
#ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_counts.{year_s}-{year_s+4}.yearly.nc'
ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_counts.TS17.nens30.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
#err = cim(da-ctl, dim='en')/meanCtl * 100
err = (da-ctl).std(dim='en')/da.en.size**0.5/meanCtl * 100
mean = da.mean('en')
anom = (mean - meanCtl)/meanCtl * 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/{volc}_PI_ens_noleap/netcdf/tc_counts.{year_s}-{year_s+4}.yearly.nc'
ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_counts.TS17.nens30.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
#err = cim(da-ctl, dim='en')/meanCtl * 100
err = (da-ctl).std(dim='en')/da.en.size**0.5/meanCtl * 100
mean = da.mean('en')
anom = (mean - meanCtl)/meanCtl * 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/{volc}_PI_ens_noleap/netcdf/tc_counts.{year_s}-{year_s+4}.yearly.nc'
ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_counts.TS17.nens30.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
#err = cim(da-ctl, dim='en')/meanCtl * 100
err = (da-ctl).std(dim='en')/da.en.size**0.5/meanCtl * 100
mean = da.mean('en')
anom = (mean - meanCtl)/meanCtl * 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/{volc}_PI_ens_noleap/netcdf/tc_counts.{year_s}-{year_s+4}.yearly.nc'
ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_counts.TS17.nens30.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
#err = cim(da-ctl, dim='en')/meanCtl * 100
err = (da-ctl).std(dim='en')/da.en.size**0.5/meanCtl * 100
mean = da.mean('en')
anom = (mean - meanCtl)/meanCtl * 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/{volc}_PI_ens_noleap/netcdf/tc_counts.{year_s}-{year_s+4}.yearly.nc'
ifile = f'/tigress/wenchang/analysis/TC/{volc}_PI_ens_noleap/netcdf/tc_counts.TS17.nens30.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
#err = cim(da-ctl, dim='en')/meanCtl * 100
err = (da-ctl).std(dim='en')/da.en.size**0.5/meanCtl * 100
mean = da.mean('en')
anom = (mean - meanCtl)/meanCtl * 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')
#plt.show()

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

if __name__ == '__main__':
    tt.check('**done**')
    plt.show()