#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Thu Oct  3 15:49:44 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 geoplots import yticks2lat, xticksmonth
 
# Parameters
cmip = 'CMIP5'#'CMIP6'
print(cmip)
figname = f'fig.aod550nm.{cmip}.png'
data_name = 'extsw_b07'
if cmip == 'CMIP6':
    ncfile = '/tigress/gvecchi/CM2.5/input/VOLCANIC/CMIP6/extsw_V3_DATATROP_RCP.nc'
elif cmip == 'CMIP5':
    ncfile = '/tigress/gvecchi/CM2.5/input/VOLCANIC/CMIP5/extsw_data.nc'

def ext2aod(ds, band=7):
    '''calculate volcanic aerosol AOD based on extinction coefficient.'''
    Rd = 287#J/K/kg
    T = 300#K
    g = 9.8#m/s**2
    data_name = f'extsw_b{band:02d}'
    if band == 7:
        new_name = 'AOD_550nm'
    else:
        new_name = f'AOD_b{band:02d}'
    da = ds[data_name]
    z = da.pfull.pipe(lambda x: x.isel(pfull=0)/x) \
        .pipe(lambda x: np.log(x))*Rd*T/g
    da = da.rename({'pfull': 'z'}).assign_coords(z=z.values).integrate('z') \
        .rename(new_name)
    
    return da

# data
ds = xr.open_dataset(ncfile)
da = ext2aod(ds, band=7).mean('lon')

# fig, CMIP6 AOD at 550nm
scale_factor = 1
da_ = da.pipe(lambda x:x*scale_factor)
da_['time'] = da_.indexes['time'].to_datetimeindex()
bbox = {'facecolor': 'w', 'alpha':.5, 'boxstyle': 'round'}

fig, axes = plt.subplots(3,1,figsize=(8,7), sharey=True)
cmap = 'OrRd'
levels = np.arange(0,0.46,0.05)

ax = axes[0]
# plt.sca(ax)
da_.sel(time=slice('1850', '1899')).plot.contourf(ax=ax, x='time', cmap=cmap, levels=levels,
                                                 add_colorbar=False)

ax.set_title(cmip)
ax.set_xlabel('')
ax.set_ylabel('lat')
ax.grid(True)
ax.axhline(0, color='gray', ls='--')
ax.text('1886-06', -6, 'Krakatoa\nIndonesia\nAug. 1883', va='center', bbox=bbox)

ax = axes[1]
# plt.sca(ax)
da_.sel(time=slice('1900', '1949')).plot.contourf(ax=ax, x='time', cmap=cmap, levels=levels,
                                                 add_colorbar=False)
ax.set_title('')
ax.set_xlabel('')
ax.set_ylabel('lat')
ax.grid(True)
ax.axhline(0, color='gray', ls='--')
ax.text('1915-01', 58, 'Novarupta\nAlaska\nJun. 1912', va='center', bbox=bbox)
ax.text('1905-01', 15, 'Santa Maria\nGuatemala\nOct. 1902', va='center', bbox=bbox)


ax = axes[2]
im = da_.sel(time=slice('1950', '1999')).plot.contourf(ax=ax, x='time', cmap=cmap, levels=levels,
                                                      add_colorbar=False)
ax.set_title('')
ax.grid(True)
ax.set_ylabel('lat')
ax.axhline(0, color='gray', ls='--')
ax.text('1994-06', 15, 'Pinatubo\nPhilippine\nJun. 1991', va='center', bbox=bbox)
ax.text('1985-01', 17, 'Chichon\nMexico\nMar. 1982', va='center', bbox=bbox)
ax.text('1966-01', -8, 'Agung\nIndonesia\nMar. 1963', va='center', bbox=bbox)


yticks2lat(np.arange(-60,61,30))

ax = fig.add_axes((.23,.08,.6,.02))
plt.colorbar(im, cax=ax, orientation='horizontal', label=f'Volcanic Aerosol AOD at 550nm', extend='both')

for ax in axes:
    plt.setp(ax.get_xticklabels(), visible=True, rotation=30, horizontalalignment='right')

plt.tight_layout(rect=(0,.1,1,1))
if figname is not None:
    plt.savefig(figname)

if __name__ == '__main__' and False:
    tformat = '%Y-%m-%d %H:%M:%S'
    t0 = datetime.now()
    print('[start]:', t0.strftime(tformat))
    
    t1 = datetime.now()
    print('[end]:', t1.strftime(tformat))
    print('[total time]:', f'{(t1-t0).seconds:,} seconds')