"""
This file is called by run_pde.sh, it generates the PDE data. First the diffusion mapping module is called
this maps tree-density values to velocity values and produces a map of the uk where each point correponds
to the wave-speed at that point of a particular pathogen. The mapping funciton is generated by ensemble
results from the SSTML.
"""
import sys, os
PATH_model = os.getcwd()+'/model'
sys.path.append(PATH_model)
PATH_diff = os.getcwd()+'/diffusion_mapping'
sys.path.append(PATH_diff)
import numpy as np
import time
import matplotlib.pyplot as plt
import pde_model
import diff_map_gen


in_arr = sys.argv  # System Inputs
date_time, current_time, hpc_mode = in_arr[1:]
ell = {1: "25", 2: "50", 3: "75", 4: "100"}
ell_ = ell[4]
domain_name = 'Fex-cg-1'
vmap_name = "R0_10_ell_{}".format(ell_)
sim_name = "test-Fex-R0_10-L_{}m".format(ell_)

# ---- Model Parameters ---- #
# epi center : infected at T=0
# plt_epi : plot check epi before running
# partial : if True, partial sub-section of data, [L_row, H_row, L_col, H_col]
# vmap : mapping function velocity to tree density
# domain_name : name of tree data \in [Oak, Ashe,..., Willow]
# ell_ : dispersal distance (m)
# R0 : Basic reproduction ratio (day^-1)

model_params = {"epi_c": [695, 700, 555, 560], "plt_epi": False, "partial": [True,  [800, 900, 300, 400]],
                "vmap": vmap_name, "domain_name": domain_name, "dispersal": ell_, "R0": 10, "sim_name": sim_name}

# ---- FD_solver_settings ---- #
# tend : time-steps iterated in days
# dx, dy : spatial descreteization parameters in units km^2
# CFL_max : condition for numerical stability
# N_saves : number of saves in simulation
# sim_name : output name
FD_solver_settings = {"tend": 200, "dx": 1, "dy": 1, "CFL_max": 0.50, "N_naves": 100, "sim_name": sim_name}


print("Running PDE : {}".format(vmap_name))
print("--> save dir :  {} ".format(sim_name))
# GENERATE diffusion map based on input of L, beta and a domain (in this case a map of abundance)
diffusion_map, species_number_map, growth_map, sea_map, uk = diff_map_gen.main(model_params, plt_figs=False)

sys.exit()

# SET growth to constant value of 1 (non-dimensionalized)
# maps_ : this is fed into pde_model
maps_ = np.zeros(shape=(5, growth_map.shape[0], growth_map.shape[1]))
maps_[0], maps_[1], maps_[2], maps_[3], maps_[4] = diffusion_map, growth_map, species_number_map, sea_map, uk

params["dim"] = diffusion_map.shape
# - epi_center : point of disease introduction
#   port of Immingham ~= [560, 570, 455, 465]
t_0 = time.time()
inf_tseries = pde_model.main(params, maps_)
t_1 = time.time()
t = (t_1 - t_0) / 60
np.save(sim_name + '_response_c', inf_tseries)
print(sim_name + '_response_c : saved...')
# Print time elapsed in minutes.
print("Time elapsed: ", t, ' (mins)')
plt.show()


