#######################################################################################################
# Create a `pomp` model to use the climate data gathered with the `nasapower` and using pre-defined parameters
# defined externally of the function. 
#######################################################################################################
Create_Model_Final_TWO_Deltas <- function(covars_df, lin_bool_val = T) {
  
  # Args: 
  #   covars_df: 
  #     - Number of the week `week_no`
  #     - Time series of normalised temperature measures
  #     - Time series of normalised relative humidity measures
  #     - Binary variable for the effects of school or holiday `school_F`
  #     - school holiday calendar
  #     - value of imported infects `import`
  #   lin_bool_val: should the transmission term be linearised? (boolean)
  # Returns: POMP model 
  
  # Check that covariates are present and formatted
  stopifnot(all(c("week_no", "Te_norm_Now", "RH_norm_Now", "Te_norm_Lag", "RH_norm_Lag", "T_school", "school_F") %in% colnames(covars_df)))
  stopifnot(min(covars_df$week_no) == 0)
  covars_df = arrange(covars_df, week_no)
  
  
  # Csnippet for stochastic model ----
  stoch_process_model <- Csnippet("
  
  double S_new, I_new, R_new, CC_new;
  
  // Calculate transmission rate
  double Beta = R0 * (mu + 1.0); // Mean transmission rate
  
  // Adding the climate-induced seasonality from the covariate table
  Beta *= exp(e_Te_Now * Te_norm_Now + e_RH_Now * RH_norm_Now + e_Te_Lag * Te_norm_Lag + e_RH_Lag * RH_norm_Lag); // Add the seasonal beta component
  
  // Environmental stochasticity -> defined in the parameters
  double dW = rgammawn(sigma_beta, 1.0); // Mean: 1, SD: sigma_beta
  Beta *= dW; 
  
  // Term-Time forcing. This is applied iff `Term_Time == 1`
  Beta *= (1 + (Term_Time * (school_F - 1)));

  // Force of infection
  double lambda_t = Beta * ((I + import) / N); // Force of infection
  double p_SI = lin_bool ? lambda_t : (1.0 - exp(-lambda_t)); // Proportion of infected during time step

  // Difference equations
  S_new = S + mu * N + (1.0 - eps) * I + alpha * R - (p_SI + mu) * S; 
  I_new = p_SI * S - mu * I;  
  R_new = R + eps * I - (alpha + mu) * R;
  CC_new = p_SI * S;
  
  // Enforce non-negative and bounded populations
  S_new = fmax(fmin(S_new, N), 0.0);
  I_new = fmax(fmin(I_new, N), 0.0);
  R_new = fmax(fmin(R_new, N), 0.0);
  CC_new = fmax(fmin(CC_new, N), 0.0);
  
  // Check and update
  S = S_new;
  I = I_new;
  R = R_new;
  CC = CC_new;
  ")
  
  # Csnippet to generate observations and estimations
  robs_model <- Csnippet("
  CC_obs = rnbinom_mu( 1.0 / rho_k, rho_mean * CC); // args: size, mean
")
  
  # Csnippet for the initial conditions of the states variables
  # All variables are initialized at the endemic equilibrium of the seasonally-unforced model
  init_mod <- Csnippet("
  double S_star = N / R0;
  S = S_star;
  I = (alpha + mu) / (alpha + mu + eps) * (N - S_star);
  R = eps / (alpha + mu + eps) * (N - S_star);
  CC = 0;
")
  
  
  # Create pomp object
  mod <- pomp(data = data.frame(week = 1:(max(covars_df$week_no)), CC_obs = NA),
              times = "week",
              t0 = 0,
              obsnames = "CC_obs",
              covar = covariate_table(select(covars_df, c("week_no", "Te_norm_Now", "RH_norm_Now",
                                                          "Te_norm_Lag","RH_norm_Lag", "T_school",
                                                          "school_F")),
                                      times = "week_no", order = "constant"),
              statenames = c("S", "I", "R", "CC"),
              rprocess =  discrete_time(step.fun = stoch_process_model, delta.t = 1),
              rmeasure = robs_model,
              rinit = init_mod,
              params = c("mu" = 1 / 80 / 52, # Birth rate (per week)
                         "N" = 5e6, # Total population size
                         "R0" = 1, # Basic reproduction number
                         "sigma_beta" = 0, # SD of environmental noise
                         "eps" = 1, # Fraction of infections conferring immunity
                         "alpha" = 0, # Waning immunity rate (per week)
                         "rho_mean" = 1, # Average reporting probability
                         "rho_k" = 0.04, # Reporting over-dispersion
                         "Term_Time" = 1, # Apply Term-Time forcing or not
                         "lin_bool" = as.integer(lin_bool_val), # Choose if the transmission term has to be linearised
                         "e_Te_Now" = 0, # Effect of temperature on transmission
                         "e_RH_Now" = 0, # Effect of relative humidity on transmission
                         "e_Te_Lag" = 0,
                         "e_RH_Lag" = 0,
                         "import" = 0), # Number of imported infections
              paramnames = c("mu", "N", "R0", "sigma_beta", "eps", "alpha", 
                             "rho_mean", "rho_k", "Term_Time", "lin_bool",
                             "e_Te_Now", "e_RH_Now", "e_Te_Lag", "e_RH_Lag",
                             "import"), 
              partrans = parameter_trans(log = c("R0", "alpha", "rho_k"), 
                                         logit = c("rho_mean"))
  )
  return(mod)
}