#!/usr/bin/env python
import xarray as xr
import matplotlib.pyplot as plt
import xfilter

basin = 'global'#,'NA'
ylabel = f'# of {basin} TCs'
figname = f'fig_TC_counts__{basin}.png'
alpha = 0.5
thin = 0.5
thick = 2
n_years = 9
plt.figure(figsize=(8,4))

# ctl
ifile = '/tigress/wenchang/analysis/TC/CTL1860_newdiag_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0001-1000.yearly.nc'
ds = xr.open_dataset(ifile)
da = ds[basin]
da.plot(color='k', lw=thin, alpha=alpha)
da.filter.lowpass(1/n_years).plot(color='k', lw=thick, label='CTL')

# m6
ifile = '/tigress/wenchang/analysis/TC/m6p0sol_CTL1860_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0101-0597.yearly.nc'
ds_m6 = xr.open_dataset(ifile)
da = ds_m6[basin]
da.plot(color='C0', lw=thin, alpha=alpha)
da.filter.lowpass(1/n_years).plot(color='C0', lw=thick, label='$-6$% solar')

# p6
ifile = '/tigress/wenchang/analysis/TC/p6p0sol_CTL1860_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0101-0400.yearly.nc'
ds_p6 = xr.open_dataset(ifile)
da = ds_p6[basin]
da.plot(color='C1', lw=thin, alpha=alpha)
da.filter.lowpass(1/n_years).plot(color='C1', lw=thick, label='+6% solar')

# m6 recover
ifile = '/tigress/wenchang/analysis/TC/m6p0sol_CTL1860_recover0501_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0501-0900.yearly.nc'
ds_m6r = xr.open_dataset(ifile)
da = ds_m6r[basin]
da.plot(color='C2', lw=thin, alpha=alpha)
da.filter.lowpass(1/n_years).plot(color='C2', lw=thick, label='$-$6% solar recover')


# m1
color = 'C3'
label = '$-1$% solar'
ifile = '/tigress/wenchang/analysis/TC/m1p0sol_CTL1860_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0101-0600.yearly.nc'
ds_m6 = xr.open_dataset(ifile)
da = ds_m6[basin]
da.plot(color=color, lw=thin, alpha=alpha)
da.filter.lowpass(1/n_years).plot(color=color, lw=thick, label=label)

# p1
color = 'C4'
label = '$+1$% solar'
ifile = '/tigress/wenchang/analysis/TC/p1p0sol_CTL1860_tigercpu_intelmpi_18_576PE/netcdf/tc_counts.0101-0600.yearly.nc'
ds_m6 = xr.open_dataset(ifile)
da = ds_m6[basin]
da.plot(color=color, lw=thin, alpha=alpha)
da.filter.lowpass(1/n_years).plot(color=color, lw=thick, label=label)

plt.legend()
if ylabel is not None:
    plt.ylabel(ylabel)

plt.tight_layout()
if figname is not None:
    plt.savefig(figname, dpi=128)