#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Apr 22 21:19:03 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
#
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
regions = dict(
    London=(51.5074, 0.1278),
    NYC=(40.7128, -74.0060+360),
    LA=(34.0522, -118.2437+360),
    Wuhan=(30.5928, 114.3055),
    Delhi=(28.7041, 77.1025),
    Maracaibo=(10.6427, -71.6125+360),
    Kinshasa=(-4.4419, 15.2663),
    Jakarta=(-6.2088, 106.8456),
    Johannesburg=(-26.2041, 28.0473),
    Buenos_Aires=(-34.6037, -58.3816+360),
    Melbourne=(-37.8136, 144.9631)
    
)


from numpy import pi as Pi
def get_q(case='cos', qmin=0.002, qmax=0.014, qrange=None, phi0=None, qyear=None):
    '''obtain specific humidity annual cycle'''
#     case = 'merra2'
#     qmin, qmax = 0.002, 0.018# units: kg/kg   
    if case == 'cos':
        if phi0 is None:
            phi0 = -Pi
        if qrange is None:
            qrange = qmax - qmin
        else: #overide qmax according to qmin and qrange
            qmax = qmin + qrange
        t = np.arange(0, 365)
        t = xr.DataArray(t, dims='time', coords=[t])
        t.attrs['units'] = 'day'
        qmean = (qmin + qmax)/2
        cycle = np.cos(2*Pi*t/365 + phi0) # normalized to have range of [0, 1]
        q = qmean + (qrange/2)*cycle
        q.name = 'q'
        q.attrs['units'] = 'kg/kg'
        s = f'cos-curve q: '
        if not hasattr(qmin, 'ndim'):
            s += f'qmin = {qmin}; '
        if not hasattr(qmax, 'ndim'):
            s += f'qmax = {qmax}; '
        if not hasattr(phi0, 'ndim'):
            s += f'phi0 = {phi0/Pi:.1f}xPi; '
        q.attrs['note'] = s
    elif case == 'era5':
        ifile = '/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/clim/q2m.run7clim.2010-2019.nc'
        q = ( 
            xr.open_dataset(ifile)
            ['q2m']
            .rename(dayofyear='time')
            .load()
        )
        q.attrs['note'] = case
    elif case == 'era5_single_year':
        ifiles = f'/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/era5.2m_specific_humidity.daily.{qyear}-??.nc'
        q = ( 
            xr.open_mfdataset(ifiles, combine='by_coords')
            ['q2m']
            #.load()
        )
        q.attrs['note'] = f'{case} {qyear}'
    elif case == 'merra2':
        ifile = '/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/merra_weekly_clim.nc'
        q2m = (
    xr.open_dataset('/tigress/gvecchi/ANALYSIS/DISEASE/COVID-19/merra_weekly_clim.nc')
    .rename(WEEK_QV2M='q2m', TT='time', LAT='latitude', LON='longitude')['q2m']
)
        q = ( 
            xr.open_dataset(ifile)
            .rename(WEEK_QV2M='q2m', TT='time', LAT='latitude', LON='longitude')
            ['q2m']
            .pipe(lambda x: x.assign_coords(time=x.time.dt.dayofyear))
            .reindex(time=np.arange(1, 366), method='nearest')
            .load()
        )
        # shift from 0-center to 180E-center
        if q.longitude.min() < 0:
            q = q.roll(longitude=q.longitude.size//2)
            q = q.assign_coords(
                longitude=q.longitude.where(q.longitude>=0, other=q.longitude+360).values
            )
        q.attrs['note'] = case
        
    return q

climdepens ={
    'HKU1': dict(a=-227.5, R0_min=1.4, R0_max=2.5),
    'OC43': dict(a=-32.5, R0_min=1.4, R0_max=2.5)
}
def q2R0(q, a, R0_min, R0_max):
    '''R0 dependence on climate (specific humidity)'''
    R0 = (R0_max - R0_min) * np.exp(a*q) + R0_min
    if isinstance(R0, xr.DataArray):
        R0.name = 'R0'
    
    return R0


def R0_at_t(R0, t, cycle=True):
    '''R0 as function of t (time) and specific humidity (q).
    q has dims of (n_time, n_grid).'''
    #R0_min, R0_max = 1.4, 2.5
    #a = -227.5 for HKU1 and -32.5 for OC43
    if cycle:# cycle through the first 365 daily specific humidity
        R0t = R0.isel(time=int( np.mod(t, 365) ), drop=True)
    else:
        R0t = R0.isel(time=int(t), drop=True)
    return R0t

Ldis = {
    'HKU1': 66.25*7,
    'OC43': 62.5*7
} 
 
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()