#!/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 = 'global'
figname = f'fig.tcCounts.anom.nudge.{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: nudged exps'
alpha = 0.3
capsize = 3 
fmt = 'o'
bar_width = 0.1
ecolor = 'gray'

volc = 'CTL'
color = 'k'
ifile = '/tigress/wenchang/analysis/TC/nudgeclimo_all_model_CTL1860_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0111-0124.yearly.nc'
da = xr.open_dataset(ifile)[basin]
# transform to ensembles
das = []
for i in range(10):
    das.append(da.isel(year=slice(i, i+5)).assign_coords(year=years))
da = xr.concat(das, dim=pd.Index(range(1, 11), 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}_ens_noleap_nudgeclimo_all_model1860/netcdf/tc_counts.{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}_ens_noleap_nudgeclimo_all_model1860/netcdf/tc_counts.{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}_ens_noleap_nudgeclimo_all_model1860/netcdf/tc_counts.{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'
#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*3
ifile = f'/tigress/wenchang/analysis/TC/{volc}_ens_noleap_nudgeclimo_all_model1860/netcdf/tc_counts.{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=4, loc='best')

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