#############################################################################
# Code to simulate 100 Time Series replicates and fit
# 2 GAM formulations to each of these repetitions.
# This is done for 50 parameters sets (corresponding to the 50 different synthetic pathogens).
# Moreover, each run uses weather data from a specific location.
# This can all be run either on a HPC SLURM cluster via `master_HPC.sh`
# or locally, via `local_master.R`. For more detail information, please refer to
# the `README.md` file
#############################################################################

# Clean environment and source functions/libs ----
source("libs_scripts.r")

# First check if running on local or on cluster ----
is_cluster <- nzchar(Sys.getenv("SLURM_JOB_ID")) # if true the script is running on cluster, otherwise on local

# Define parameters tested ----
  location     <- Sys.getenv("location")
  sigma_beta   <- as.numeric(Sys.getenv("sigma_betas_val"))
  ID           <- as.integer(Sys.getenv("ID"))
  True.e_Te    <- as.numeric(Sys.getenv("True_e_Te"))
  True.e_RH    <- as.numeric(Sys.getenv("True_e_RH"))
  rho_K        <- as.numeric(Sys.getenv("rho_K"))
  rho_Mean     <- as.numeric(Sys.getenv("rho_Mean"))
  import       <- as.numeric(Sys.getenv("import_val"))
  School_Term  <- as.integer(Sys.getenv("School_Term"))

# Set parameters ----
## Setting the number of scenarios ----
pars_num <- 50

## Define number of replicates to simulate for each scenario ----
# This below depends if running on HPC cluster or locally
n_reps <- ifelse(is_cluster, 100, 5) # Define subset of replicates to be run in local environment

## Setting parameter sets ----
### Pathogend-specific parameters sampling  ----
mu_year <- 1 / 80 # Natural birth and death rates 
parameters <- pomp::sobol_design(lower = c(A = 1, alpha_year = 1/20), upper = c(A = 10, alpha_year = 1), nseq = pars_num)
# alpha_year = alpha * 52
# 20 yrs immunity -> alpha = 1 / (52 * 20) -> alpha_year = 52 * 1 / (20 * 52) = 1 / 20
# 1 yr immunity -> alpha = 1 / (52 * 1) -> alpha_year = 52 * 1 / (1 * 52) = 1
parameters <- parameters |>
  dplyr::mutate(R0 = 1 + (1 + 1 / (alpha_year + mu_year)) / (A * (mu_year + 1))) # Obtain the R0 for each combination

# Add here the weekly alpha rate 
parameters <- parameters |> dplyr::mutate(alpha = alpha_year / 52) |>
  dplyr::mutate(eps = rep(1,nrow(parameters))) |>
  dplyr::mutate(.id = seq_len(nrow(parameters)))

## Merge all parameters in a single dataframe (across all scenarios) ----
final_pars <- expand_grid(parameters, True.e_Te, True.e_RH, sigma_beta)
final_pars <- as.data.frame(final_pars)

## .id specific parameter values set to `parms` ----
parms <- c("mu" = 1 / 80 / 52, # Birth rate 
          "N" = 5E6, # Total population size
          "R0" = final_pars[ID, "R0"], # Basic reproduction number
          "sigma_beta" = final_pars[ID, "sigma_beta"], # SD of environmental noise (set to 0 for a deterministic process model)
          "eps" = final_pars[ID, "eps"], # Fraction of infections conferring sterilizing immunity (1: SIR, 0: SIS)
          "alpha" = final_pars[ID, "alpha"], # Rate of waning immunity. THIS IS ACTUALLY THE MEAN OF AN EXP DECAY.
          "rho_mean" = rho_Mean, # Average reporting probability 
          "Term_Time" = as.numeric(School_Term), # Introduce or not the term-time coming from school contacts variations
          "e_Te" = True.e_Te, # True effects of weather variables on transmission
          "e_RH" = True.e_RH, # True effects of weather variables on transmission
          "import" = import) # Prevalence of imported infections

# Reporting over-dispersion - i.e. measurement noise
parms <- append(parms, c("rho_k" = rho_K), after = which(names(parms)=="rho_mean"))

# Add final parameters needed
final_pars <- final_pars |> dplyr::mutate(N = parms["N"],
                                         rho_mean = parms["rho_mean"],
                                         Term_Time = parms["Term_Time"],
                                         rho_k = parms["rho_k"])

# Weather data ----
## Import climate data of interest ----
climate <- readRDS("weather_data/picked_Locations_weather.rds") |>
  dplyr::filter(loc == location) |> dplyr::arrange(week_no)

## Add the holiday categorical index ----
# This results in 37/52 weeks of school
# +1 for increased contacts, -1 for decreased contacts;
climate <- climate |> dplyr::mutate(T_school = case_when(
  week_in_year %in% c(1, 24:35, 52) ~ -1,
  TRUE ~ 1
))

### Add school term forcing to the climate dataframe ----
e_school <- 0.3 / 1.7 # Tuned to have 30% reduction in contacts on holiday terms
climate <- climate |> dplyr::mutate(school_F_raw = 1 + (T_school * e_school)) |> 
  dplyr::mutate(school_F = school_F_raw / mean(school_F_raw))

## Create covariate table for the pomp model with climate dataframe ----
covars <- climate |> 
  dplyr::select(week_no, T_school, school_F, Te_norm, RH_norm) |>
  dplyr::arrange(week_no)

# Creating POMP model ----
Model <- Create_Model_Final(covars_df = covars, lin_bool_val = TRUE)

# Setting parameters of the pomp model
pomp::coef(Model, names(parms)) <- unname(parms)

# Generate simulations ----
## MIND: tune local cores for increased computational power
n.cores <- if (is_cluster) {
  as.integer(Sys.getenv("SLURM_CPUS_PER_TASK"))
} else {
  max(2, min(5, parallel::detectCores() - 2))
}

if (is_cluster) {
  future::plan(future::multicore, workers = n.cores)
} else {
  future::plan(future::multisession, workers = n.cores)
}

## Set seed and initialize parallel simulations
set.seed(123) 
tmp <- foreach(i=1:n_reps, .combine='comb_rbind', .options.future = list(seed = TRUE)) %dofuture% {
  foo <- simulate(object = Model,
                 nsim = 1,
                 format = "data.frame")
  foo <- foo |> dplyr::mutate(rep = i)
  list(foo)
}
plan(sequential)

## Get the results of the simulations
sims_all <- tmp[[1]]
rm(tmp)

## Cleaning data
sims_all <- sims_all |> 
  dplyr::rename("week_no" = "week") |> 
  dplyr::mutate(.id = as.integer(ID), 
                rep = as.integer(rep))

sims_all <- sims_all |> left_join(y = climate |> dplyr::select(week_no, week_in_year, yr, matches("Te|RH"), T_school), by = c("week_no"="week_no")) |> 
  dplyr::arrange(.id, rep, week_no)

sims_all <- sims_all |> dplyr::group_by(rep) |> dplyr::arrange(week_no) |> 
  dplyr::mutate(SI_lag = dplyr::lag(x = S * I, n = 1L, order_by = week_no),
                CC_lag = dplyr::lag(x = CC, n = 1L, order_by = week_no),
                CC_obs_lag = dplyr::lag(x = CC_obs, n = 1L, order_by = week_no),
                T_school_lag = dplyr::lag(x = T_school, n = 1L, order_by = week_no))

# Fit GAMs on simulated data ----
## GAM models to test -----
type_smooth_vec <- c(
  "CC_true", "CC_stand_smooth_AC"
)

## Choose which function to be used w.r.t. term time forcing
if(School_Term == 1){chosen_function <- f_reg_Forcing}else{chosen_function <- f_reg}

## Initialise the output structure -----
sim_reg <- NULL

## Initialise parallel computation ----
if (is_cluster || .Platform$OS.type == "unix") {
  my.cluster <- parallel::makeForkCluster(length(type_smooth_vec))
} else {
  my.cluster <- parallel::makeCluster(length(type_smooth_vec))
}

doParallel::registerDoParallel(cl = my.cluster) # Registering cluster

sim_reg <- foreach(idx=1:length(type_smooth_vec)) %dopar% {
  s <- (type_smooth_vec[idx])
  sim_reg[[s]] <- sims_all |>
    dplyr::group_by(.id, rep, .add = TRUE) |> 
    dplyr::group_nest() |> 
    dplyr::mutate(reg = purrr::map(.x = data,
                                   .f = chosen_function,
                                   type_mod = s,
                                   SchoolForcing = School_Term,
                                   .progress = F)) |>
    tidyr::unnest(cols = "reg")
  
  sim_reg[[s]] <- sim_reg[[s]]  |>
    tidyr::pivot_longer(cols = e_Te:R2) |>  # This order depends on the one defined within the functions `f-Regressions*`
    dplyr::mutate(type = s) |> dplyr::select(-data)
}

parallel::stopCluster(my.cluster) 

## Clean regressions' outputs ----
sim_reg <- dplyr::bind_rows(sim_reg) |>
  tidyr::pivot_wider(names_from = "name", values_from = "value") |>
  dplyr::group_by(type) |>
  dplyr::mutate(bias_Te = abs(e_Te - True.e_Te),
                bias_RH = abs(e_RH - True.e_RH)) |>
  dplyr::mutate(TP_Te = ifelse(e_Te_low_CI < True.e_Te & e_Te_high_CI > True.e_Te, 1, 0), 
                e_Te_pow = ifelse(sign(e_Te_low_CI) == sign(True.e_Te) & sign(e_Te_high_CI) == sign(True.e_Te), 1, 0),
                TP_RH = ifelse(e_RH_low_CI < True.e_RH & e_RH_high_CI > True.e_RH, 1, 0), 
                e_RH_pow = ifelse(sign(e_RH_low_CI) == sign(True.e_RH) & sign(e_RH_high_CI) == sign(True.e_RH), 1, 0))

# Save everything with the ID ----
roiLocation <- sub(",.*", "", location)
if (!dir.exists("Results")) {
  dir.create("Results", recursive = TRUE)
}

my_path = paste("Results/R_sims_Te",final_pars[1, "True.e_Te"], "_RH", final_pars[1,"True.e_RH"],
                "_sigma",final_pars[1, "sigma_beta"], "_",
                ifelse(final_pars[1, "rho_mean"] != 0, paste("rho_mean", final_pars[1, "rho_mean"], "_", sep="") , ""),
                roiLocation, "_",
                sprintf("_rho_k%s", gsub("\\.", "-", as.character(final_pars[1, "rho_k"]))),
                sprintf("_import_%s", as.character(import)),
                ifelse(School_Term == 1, "_School-Forced", ""),
                sep = "")
if (!dir.exists(my_path)) {dir.create(my_path)}
save(sims_all,
     sim_reg,
     climate,
     final_pars,
     file=sprintf("%s/results_%s.RData", my_path, ID))