#!/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'
figname = f'fig.tcCounts.{basin}.png'
figsize = (8, 4)#None
fig = plt.figure(figsize=figsize)
years = range(1,6) 
xticks = years
ylabel = f'# of {basin} TCs'
xlabel = 'year'
alpha = 0.3
capsize = 5
fmt = 'o'

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))#year 11-45 are used to construct the control ensemble
# 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') )
err = cim(da, dim='en')
mean = da.mean('en')
plt.errorbar(da.year, mean, yerr=err, color=color, capsize=capsize, label=volc, fmt=fmt)
#plt.fill_between(da.year, mean-err, mean+err, color=color, alpha=alpha)
#mean.plot(color=color, label=volc)

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)
#for en in da.en:
#   da.sel(en=en).plot(color='C0', lw=1, alpha=0.5)
err = cim(da, dim='en')
mean = da.mean('en')
plt.errorbar(da.year+offset, mean, yerr=err, color=color, capsize=capsize, label=volc, fmt=fmt)
#plt.fill_between(da.year, mean-err, mean+err, color=color, alpha=alpha)
#mean.plot(color=color, label=volc+f'-{year_s}')

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.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
err = cim(da, dim='en')
mean = da.mean('en')
plt.errorbar(da.year+offset, mean, yerr=err, color=color, capsize=capsize, label=volc, fmt=fmt)
#plt.fill_between(da.year, mean-err, mean+err, color=color, alpha=alpha)
#mean.plot(color=color, label=volc+f'-{year_s}')

volc = 'Pinatubo'
year_s = 1991
color = 'C2'
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.{year_s}-{year_s+4}.yearly.nc'
da = xr.open_dataset(ifile)[basin].assign_coords(year=years)
err = cim(da, dim='en')
mean = da.mean('en')
plt.errorbar(da.year+offset, mean, yerr=err, color=color, capsize=capsize, label=volc, fmt=fmt)
#plt.fill_between(da.year, mean-err, mean+err, color=color, alpha=alpha)
#mean.plot(color=color, label=volc+f'-{year_s}')

#volc = 'StMaria'
#year_s = 1902
#color = 'C1'
#offset = 0.1*4
#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, dim='en')
#mean = da.mean('en')
#plt.errorbar(da.year+offset, mean, yerr=err, color=color, capsize=capsize, label=volc, fmt=fmt)

if xticks is not None:
    plt.xticks(xticks)
if ylabel is not None:
    plt.ylabel(ylabel)
if xlabel is not None:
    plt.xlabel(xlabel)
#plt.legend(ncol=4, bbox_to_anchor=(0, 1), loc='lower left')
plt.legend(ncol=4, loc='upper center')
plt.show()

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