#############################################################################
# Code for the sensitivity analysis on the spline dimension.
# This script leverages on previously simulated time series
# from the baseline parameter set.
#############################################################################

# Clean environment and source functions and load libs ----
rm(list=ls())
source("libs_scripts.R")

# Get cluster environmental variables i.e.: parameters tested ----
Locations = c("Dubai", "Rio_de_Janeiro", "Rome", "Toronto")
True_e_Te = -0.04
True_e_RH = -0.04
rho_mean = 0.5
School_Term = F
years = 10
rho_K = 0.1
import <- 1e-05
sigma_val <- 0

sims_total = NULL
# First import all simulated TS neede to run the GAMs ----
for(roiLocation in Locations){
  data_Path <- paste("Results_Aggregated/R_sims_Te",
                        True_e_Te, "_RH", True_e_RH,
                        "_",
                        "sigma", sigma_val,
                        "_",
                        "rho_mean", rho_mean,
                        "_",
                        roiLocation, "_",
                        ifelse(School_Term == T, "_School-Forced", ""),
                        sprintf("_rho_k%s", gsub("\\.", "-", as.character(rho_K))),
                        sprintf("_import_%s", as.character(import)),
                        sep = "")
  final_pars <- readRDS(paste(data_Path, "/final_pars.rds", sep = "")) |> dplyr::select(A,R0,alpha,.id, alpha_year)
  sims_tot <- readRDS(paste(data_Path, "/sims_tot.rds", sep = "")) |> 
    dplyr::mutate(location = roiLocation) |> dplyr::ungroup()
  
  sims_total <- sims_total |>
    dplyr::bind_rows(sims_tot)
}

# Running all fit of GAMs on simulated data ----
model <- "CC_stand_smooth_AC"

Conds <- expand.grid(
                    maximum_knots = 75,
                    scenario = unique(final_pars$.id),
                    replicate = 1:100
                    )

final_regs <- NULL
for(tmp_location in c("Rome","Rio_de_Janeiro","Dubai","Toronto")){
loc_sims <- sims_total |> dplyr::filter(location == tmp_location)
# Set up parallel computation for the conditions to be tested ----
## Initialise parallel computation ----
n.cores = as.integer(10)
future::plan(future::multisession, workers = n.cores)
set.seed(4L)
tmp_df <- foreach(idx=1:nrow(Conds), .combine = "comb_rbind", .options.future = list(seed = TRUE)) %dofuture% {
    foo <- gam(formula = CC_obs ~ 1 + Te_norm_lag + RH_norm_lag + s(week_no, k = Conds[idx,"maximum_knots"]) + log1p(CC_obs_lag),
                family = nb(link = "log"),
                data = loc_sims |> dplyr::filter( 
                                    .id == Conds[idx,"scenario"],
                                    rep == Conds[idx,"replicate"]
                                    ),
                method = "REML")
    coefs <- coef(foo)
    summ <- summary(foo)
    k_index <- mgcv::k.check(foo)[,"k-index"]
    res <- data.frame(e_Te = coefs["Te_norm_lag"], e_RH = coefs["RH_norm_lag"],
            e_Te_se = summ$se["Te_norm_lag"], e_RH_se = summ$se["RH_norm_lag"],
            edf = summ["edf"],
            k_index = k_index) |> 
        dplyr::bind_cols(.id = Conds[idx,"scenario"], rep = Conds[idx,"replicate"], k_max = Conds[idx,"maximum_knots"])
    rownames(res) <- NULL

    list(res)
}
plan(sequential)

final_res <- tmp_df[[1]]
final_regs <- rbind(final_regs,final_res |> dplyr::bind_cols(loc = tmp_location))
cat(" Done for",tmp_location)
}

if (!file.exists("Results_Aggregated/Regressions_75Knots.rds")) {
  saveRDS(final_regs, "Results_Aggregated/Regressions_75Knots.rds")
} else {
  message("File already present — not overwritten.")
}