#############################################################################
# Function with the GAM regression models to be fit to the simulations.
# Then useful outputs are ordered and outputted from the entire funtion.
#############################################################################
f_reg <- function(df, type_mod, SchoolForcing) { 

  stopifnot(SchoolForcing == F) # These GAMs do NOT comprehend Term Time forcing as covariates
  k_max <- 52
  
  if(type_mod == "CC_true") {
    M_nb <- gam(formula = CC_obs ~ 1 + Te_norm_lag + RH_norm_lag + log1p(SI_lag), 
                family = nb(link = "log"), 
                data = df, method = "REML")
  } else if(type_mod == "CC_stand_smooth_AC") {
    M_nb <- gam(formula = CC_obs ~ 1 + Te_norm_lag + RH_norm_lag + s(week_no, k = k_max) + log1p(CC_obs_lag),
                family = nb(link = "log"),
                data = df,  method = "REML")
  }
  
  # Extracting the regression coefficients
  M_est <- coef(M_nb)
  gam_summary <- summary(M_nb)
  M_est_se <- gam_summary$se
  R2 <- gam_summary$r.sq
  
  # Return
  out <- data.frame(e_Te = M_est["Te_norm_lag"], # Point estimate
                   e_Te_se = M_est_se["Te_norm_lag"], # SE of estimate
                   e_Te_low_CI = M_est["Te_norm_lag"] - (1.96 * M_est_se["Te_norm_lag"]),
                   e_Te_high_CI = M_est["Te_norm_lag"] + (1.96 * M_est_se["Te_norm_lag"]),
                   e_RH = M_est["RH_norm_lag"],
                   e_RH_se = M_est_se["RH_norm_lag"],
                   e_RH_low_CI = M_est["RH_norm_lag"] - (1.96 * M_est_se["RH_norm_lag"]),
                   e_RH_high_CI = M_est["RH_norm_lag"] + (1.96 * M_est_se["RH_norm_lag"]),
                   log_SI_lag_VAR = var(log(pmax(df[!(is.na(df$SI_lag)),]$SI_lag, 1e-6))), # Removing NAs since they're lagged cols
                   mean_log_SI_lag = mean(log(pmax(df[!(is.na(df$SI_lag)),]$SI_lag, 1e-6))),
                   R2 = R2) |>
    select(e_Te, e_Te_se, e_Te_low_CI, e_Te_high_CI, e_RH, e_RH_se, e_RH_low_CI, e_RH_high_CI, log_SI_lag_VAR, mean_log_SI_lag, R2)
  
  return(out)
}