#!/usr/bin/env python
# Wenchang Yang (wenchang@princeton.edu)
# Sun Apr 12 01:07:17 EDT 2020
if __name__ == '__main__':
    from misc.timer import Timer
    t = Timer()
import sys, os.path, os
import xarray as xr, numpy as np, pandas as pd
#import matplotlib.pyplot as plt
#more imports
#from scipy.integrate import odeint
from climSIRS import run_climSIRS
from shared import get_q, q2R0, climdepens, Ldis
#
if __name__ == '__main__':
    t.check('import done')   

#start from here
qyears = range(2019, 1978, -1)
dis = 'HKU1' #'OC43'
L = Ldis[dis]
D = 5
ii0 = 1/8e6 # initial I/N: one person in 8 million people
T = 365*10
dt = 1

if __name__ == '__main__':
    for qyear in qyears:
        t_loop = Timer(f'loop q{qyear} start')
        print(f'{qyear}q')
        ofile = f'climSIRS.{dis}.q{qyear}.nc'
        if os.path.exists(ofile):
            print('[exists]:', ofile)
            continue
        '''
        ifiles = f'/tigress/wenchang/data/era5/analysis_wy/2m_specific_humidity/daily/*.{qyear}*.nc'
        q2m = xr.open_mfdataset(ifiles)['q2m']
        q = q2m.load()#.rename(dayofyear='time') 
        '''
        q = get_q('era5_single_year', qyear=qyear).load()
        R0 = q2R0(q, **climdepens[dis])

        t_model = Timer('climSIRS start')
        ds = run_climSIRS(ii0=ii0, T=T, dt=dt, L=L, D=D, R0=R0, dis=dis)
        t_model.check(f'climSIRS q{qyear} done')
        encoding = {
            'ss': {'dtype': 'float32', 'zlib': True, 'complevel': 1},
            'ii': {'dtype': 'float32', 'zlib': True, 'complevel': 1},
            }
        print(f'saving data to netcdf ...')
        t_save = Timer('save data to netcdf')
        ds.to_netcdf(ofile, encoding=encoding)
        print('[saved]:', ofile)
        t_save.check(f'save data to netcdf done for q{qyear}')

        t_loop.check(f'loop q{qyear} done')
        
    t.check('script end')