#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Tue Apr 14 15:42:51 EDT 2020
import sys, os.path, os, datetime
#import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
from wystart import *
#
if __name__ == '__main__':
    print()
    today = datetime.date.today()
    today_s = today.strftime('%Y-%m-%d')
    tformat = '%Y-%m-%dT%H:%M:%S'
    t_start = datetime.datetime.now()
    print('[start]:', t_start.strftime(tformat))
    
#start from here
#MERRA2 q
dis = 'OC43'
qsource = 'MERRA2'
ds_sp = xr.open_dataset('/tigress/wenchang/analysis/climate_and_disease/climSIRS/ClimSIRS.OC43.merra2wclimq.cities.nc')
ds_wy = xr.open_dataset('/tigress/wenchang/analysis/climate_and_disease/climSIRS/xClimSIRS.OC43.merra2wclimq.cities.nc')

#plot
cities = ds_sp.city.values
colors = [f'C{i}' for i in range(cities.size)]
peaks = []
peak_days = []
for city,c in zip(cities,colors):
    ds_sp.ii.sel(city=city).plot(color=c, label=f'{city} (odeint)')
    ds_wy.ii.sel(city=city).plot(color='k', ls=':', label=f'{city} (RK4)')
    da_sp = ds_sp.sel(city=city).ii.squeeze().isel(t=ds_sp.sel(city=city).ii.squeeze()==ds_sp.sel(city=city).ii.max())
    da_wy = ds_wy.sel(city=city).ii.squeeze().isel(t=ds_wy.sel(city=city).ii.squeeze()==ds_wy.sel(city=city).ii.max())
    peaks.append([da_sp.item(), da_wy.item()])
    peak_days.append([da_sp.t.astype('int').item(), da_wy.t.item()])
    plt.xlim(0, 365*2)
plt.legend()
plt.title(f'{dis} {qsource}')
figname = f'fig.ii.cities.{qsource}.png'
plt.savefig(figname)

df_peaks = pd.DataFrame(peaks, index=cities, columns=['odeint', 'RK4'])
print('I/N peak')
print(df_peaks)
print()
df_peak_days = pd.DataFrame(peak_days, index=cities, columns=['odeint', 'RK4'])
print('I/N peak day')
print(df_peak_days)
 
if __name__ == '__main__':
    t_end = datetime.datetime.now()
    print('[end]:', t_end.strftime(tformat))
    print('[total time used]:', f'{(t_end - t_start).seconds:,} seconds')
    print()
    plt.show()