#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Wed Apr 22 21:17:23 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 shared import regions, get_q, climdepens, q2R0, R0_at_t, Ldis
#
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
def model_climSIRS(ss, ii, t, L, D, R0):
    '''SIRS model right hand side from Baker et al., 2020.

    Input:
    -------
        ss: S/N
        ii: I/N
        t: time (units is 'day')
        L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
        D: infection period, 5 days
        R0 = beta*D, scalar, 1-D array or time-dependent N-D DataArray

    Return:
    --------
        tuple (dsdt, didt)
        dsdt: S/N tendency
        didt: I/N tendency
    '''
    if hasattr(R0, 'time'): # R0 is time-dependent
        _R0 = R0_at_t(R0, t)
        if _R0.ndim < 1:
            _R0 = _R0.item()
    else:
        _R0 = R0
    dsdt = (1-ss-ii)/L - _R0*ii*ss/D
    didt = _R0*ii*ss/D - ii/D

    return dsdt, didt

def rk4_climSIRS(ss, ii, t, dt, L, D, R0):
    '''Step change from classic Runge–Kutta (RK4) method for SIRS model of Baker et al., 2020.

    Input:
    -------
        ss: S/N
        ii: I/N
        t: time (units is 'day')
        dt: time step
        L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
        D: infection period, 5 days
        R0 = beta*D, scalar, 1-D array or N-D time-dependent DataArray

    Return:
    --------
        tuple (dss, dii)
        dss: S/N step increase
        dii: I/N step increase
    '''
    #model_climSIRS = model_climSIRSwy
    k1s, k1i = model_climSIRS(ss=ss, ii=ii, t=t, L=L, D=D, R0=R0)
    k2s, k2i = model_climSIRS(ss=ss+k1s*dt/2, ii=ii+k1i*dt/2, t=t+dt/2, L=L, D=D, R0=R0)
    k3s, k3i = model_climSIRS(ss=ss+k2s*dt/2, ii=ii+k2i*dt/2, t=t+dt/2, L=L, D=D, R0=R0)
    k4s, k4i = model_climSIRS(ss=ss+k3s*dt, ii=ii+k3i*dt, t=t+dt, L=L, D=D, R0=R0)

    dss = dt*(k1s + 2*k2s + 2*k3s + k4s)/6
    dii = dt*(k1i + 2*k2i + 2*k3i + k4i)/6

    return dss, dii

def run_climSIRS(ii0=1e-5, ss0=None, T=1*365, dt=1, tvec=None, L=None, D=5, R0=2, dis='HKU1'):
    '''Integrate SIRS model from Baker et al., 2020.

    Input:
    -------
        I0: initial I/N
        T: time stop (start is 1; units is 'day')
        dt: time step
        tvec: will override T and dt if not None;
        L: immunity length, 7x66.25(62.5) days for HKU1 (OC43)
        D: infection period, 5 days
        R0: beta*D, scalar, 1-D array or time-dependent n-D DataArray
        dis: disease name, e.g. "HKU1" or "OC43"
        
    Return:
    --------
        ds: solution for the climSIRSwy model.
    '''
    
    if tvec is None:
        t_vec = np.arange(0, T, dt)
    else:
        t_vec = tvec
    
    if L is None and dis is not None:
        L = Ldis[dis]
    if hasattr(R0, 'time'): # R0 time-dependent
        zeros = ii0*L*D*R0.isel(time=0, drop=True)*0 # broadcast array by R0.dims[1:]
    else: # R0 time-independent
        zeros = ii0*L*D*R0*0
    # default initial S/N value determined by initial I/N value
    if ss0 is None:
        ss0 = 1 - ii0
    # broadcast initial values by adding zeros
    ii = ii0 + zeros
    ss = ss0 + zeros
    
    ss_vec = [ss,]
    ii_vec = [ii,]
    print('time step (units: day): ', end='')
    for t in t_vec[1:]:
        if np.mod(t, 365) == 0:
            print(f'365x{t//365}', end='; ')
        dss, dii = rk4_climSIRS(ss=ss, ii=ii, t=t, dt=dt, L=L, D=D, R0=R0)
        ss = ss + dss
        ii = ii + dii
        ss_vec.append(ss)
        ii_vec.append(ii)
    print()
    if isinstance(ss, xr.DataArray):
        try:
            da_ss = xr.concat(ss_vec, dim=pd.Index(t_vec, name='t'))
            da_ii = xr.concat(ii_vec, dim=pd.Index(t_vec, name='t'))
        except MemoryError:
            ss_vec = [_ss.expand_dims('dum').chunk({'dum':1}).squeeze() for _ss in ss_vec]
            ii_vec = [_ii.expand_dims('dum').chunk({'dum':1}).squeeze() for _ii in ii_vec]
            da_ss = xr.concat(ss_vec, dim=pd.Index(t_vec, name='t'))
            da_ii = xr.concat(ii_vec, dim=pd.Index(t_vec, name='t'))
    else:
        da_ss = xr.DataArray(ss_vec, dims=['t'], coords=[t_vec])
        da_ii = xr.DataArray(ii_vec, dims=['t'], coords=[t_vec])
    da_ss.attrs['long_name'] = 'S/N'
    da_ii.attrs['long_name'] = 'I/N'
    ds = xr.Dataset(dict(ss=da_ss, ii=da_ii, year=da_ss.t/365))
    ds.t.attrs['units'] = 'day'
    

    return ds 
 
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()