# Load models
MXL_wtp_Control_Not_Pred_model <- apollo_loadModel("Estimation_results/mxl/prediction/Split_samples/MXL_wtp_Control_Not_Pred")
MXL_wtp_Control_Pred_model <- apollo_loadModel("Estimation_results/mxl/prediction/Split_samples/MXL_wtp_Control_Pred")
MXL_wtp_Opt_Not_Pred_model <- apollo_loadModel("Estimation_results/mxl/prediction/Split_samples/MXL_wtp_Opt_Not_Pred")
MXL_wtp_Opt_Pred_model <- apollo_loadModel("Estimation_results/mxl/prediction/Split_samples/MXL_wtp_Opt_Pred")
MXL_wtp_Treated_Not_Pred_model <- apollo_loadModel("Estimation_results/mxl/prediction/Split_samples/MXL_wtp_Treated_Not_Pred")
MXL_wtp_Treated_Pred_model <- apollo_loadModel("Estimation_results/mxl/prediction/Split_samples/MXL_wtp_Treated_Pred")


mxl_tr <- apollo_loadModel("Estimation_results/mxl/Split_samples/MXL_wtp Treated A")
mxl_vol_tr <- apollo_loadModel("Estimation_results/mxl/Split_samples/MXL_wtp Vol_Treated A")
mxl_not_tr <- apollo_loadModel("Estimation_results/mxl/Split_samples/MXL_wtp Not_Treated A")

# List of models
models <- list(
  "Treated Not Pred" = MXL_wtp_Treated_Not_Pred_model,
  "Control Not Pred" = MXL_wtp_Control_Not_Pred_model,
  "Control Pred" = MXL_wtp_Control_Pred_model,
  "Opt Not Pred" = MXL_wtp_Opt_Not_Pred_model,
  "Opt Pred" = MXL_wtp_Opt_Pred_model,
  "Treated Pred" = MXL_wtp_Treated_Pred_model,
  "Control Split Sample" = mxl_not_tr,
  "Opt Split Sample" = mxl_vol_tr,
  "Treated Split Sample" = mxl_tr
)

group_color_mapping <- c(
  "Control" = "grey",
  "Treated" = "red",
  "Optional" = "orange"
)

# Define line type mapping for prediction status
prediction_line_mapping <- c(
  "Not Predicted" = "dashed",
  "Predicted" = "solid",
  "Split Sample" = "dotted"
)

# Create a mapping for the groups
group_mapping <- c(
  "Control Not Pred" = "Control",
  "Control Pred" = "Control",
  "Treated Not Pred" = "Treated",
  "Treated Pred" = "Treated",
  "Opt Not Pred" = "Optional",
  "Opt Pred" = "Optional",
  "Control Split Sample" = "Control",
  "Opt Split Sample" = "Optional",
  "Treated Split Sample" = "Treated"
)

# Create a mapping for prediction status
prediction_mapping <- c(
  "Control Not Pred" = "Not Predicted",
  "Control Pred" = "Predicted",
  "Control Split Sample" = "Split Sample",
  "Treated Not Pred" = "Not Predicted",
  "Treated Pred" = "Predicted",
  "Treated Split Sample" = "Split Sample",
  "Opt Not Pred" = "Not Predicted",
  "Opt Pred" = "Predicted",
  "Opt Split Sample" = "Split Sample"
)

# Define the x-range for plotting (1000 values)
x_range <- seq(-20, 100, length.out = 1000)

# Define an empty data frame to store the results
plot_data <- data.frame()

# Loop through the models and calculate the density for each
for (i in seq_along(models)) {
  model_name <- names(models)[i]
  model <- models[[i]]
  
  # Extract mean and standard deviation for 'mu_natural' and 'sig_natural'
  coef <- model$estimate
  mean <- coef["mu_natural"]
  sd <- coef["sig_natural"]
  
  # Compute the normal distribution based on the extracted mean and sd
  y <- dnorm(x_range, mean, sd)
  
  # Create a data frame for this model and append it to the main data frame
  temp_data <- data.frame(
    x = x_range, 
    density = y, 
    model = model_name,
    Group = group_mapping[[model_name]],
    Prediction = prediction_mapping[[model_name]]
  )
  
  plot_data <- bind_rows(plot_data, temp_data)
}

# Plot using ggplot2 with custom color and linetype mappings
ggplot(plot_data, aes(x = x, y = density, color = Group, linetype = Prediction)) +
  geom_line(size = 1) + 
  scale_color_manual(values = group_color_mapping) +  # Set custom colors for groups
  scale_linetype_manual(values = prediction_line_mapping) +  # Set custom linetypes for prediction status
  labs(
    title = "Normal Distributions from Multiple MXL Models",
    x = "WTP Naturalness (€/month)",
    y = "Density",
    color = "Group",
    linetype = "Prediction"
  ) +
  theme_minimal() +
  theme(legend.position = "right")


#### Z-test if distributions are different

z_value <- (MXL_wtp_Treated_Pred_model$estimate["mu_natural"] - MXL_wtp_Control_Pred_model$estimate["mu_natural"])/
  (sqrt((MXL_wtp_Treated_Pred_model$estimate["sig_natural"]/sqrt(MXL_wtp_Treated_Pred_model$nObs)) + 
     (MXL_wtp_Control_Pred_model$estimate["sig_natural"]/sqrt(MXL_wtp_Control_Pred_model$nObs))))


# Function to calculate z-value using your exact formula
calculate_z_value <- function(model_1, model_2) {
  # Calculate the z-value using your specific formula
  z_value <- (model_1$estimate["mu_natural"] - model_2$estimate["mu_natural"]) /
    (sqrt((model_1$estimate["sig_natural"] / sqrt(model_1$nObs/10)) + 
       (model_2$estimate["sig_natural"] / sqrt(model_2$nObs/10))))
  
  return(z_value)
}

# Example of using the function
z_value <- calculate_z_value(MXL_wtp_Treated_Pred_model, MXL_wtp_Control_Pred_model)

# Print the z-value
z_value


#### Plot underneath

ggplot(plot_data, aes(x = x, y = density, color = Group, linetype = Prediction)) +
  geom_line(size = 1) + 
  scale_color_manual(values = group_color_mapping) +  # Set custom colors for groups
  scale_linetype_manual(values = prediction_line_mapping) +  # Set custom linetypes for prediction status
  facet_grid(Prediction ~ ., scales = "free_y") +  # Stacks "Not Predicted" on top of "Predicted"
  labs(
    title = "Normal Distributions from Multiple MXL Models",
    x = "WTP Naturalness (€/month)",
    y = "Density",
    color = "Group",
    linetype = "Prediction"
  ) +
  theme_minimal() +
  theme(
    legend.position = "right",
    strip.text.y = element_text(size = 12, face = "bold"),  # Adjust facet labels
    panel.spacing = unit(1, "lines")  # Adds space between the two facets
  ) +
  xlim(-25, 75)


#### Compare means
# Create an empty data frame
# Function to compute CI dynamically
calculate_ci <- function(mean, se, confidence = 0.90) {
  z_value <- qnorm((1 + confidence) / 2)  # Compute Z-score for the given confidence level
  lower_ci <- mean - z_value * se
  upper_ci <- mean + z_value * se
  return(c(lower_ci, upper_ci))
}

ci_data <- data.frame()

# Loop through each model and extract mean and confidence interval
for (i in seq_along(models)) {
  model_name <- names(models)[i]
  model <- models[[i]]
  
  # Extract mean estimate and robust standard error
  mean_value <- model$estimate["mu_natural"]
  se_robust <- model$robse["mu_natural"]
  
  # Compute 90% confidence intervals dynamically
  ci_values <- calculate_ci(mean_value, se_robust, confidence = 0.90)
  
  # Append to data frame
  temp_data <- data.frame(
    Model = model_name,
    Mean = mean_value,
    Lower_CI = ci_values[1],
    Upper_CI = ci_values[2],
    Group = group_mapping[[model_name]], 
    Prediction = prediction_mapping[[model_name]]
  )
  
  ci_data <- bind_rows(ci_data, temp_data)
}

# Ensure correct order of Groups (Treated > Optional > Control)
ci_data$Group <- factor(ci_data$Group, levels = c("Treated", "Optional", "Control"))
ci_data$Model <- factor(ci_data$Model, levels = unique(ci_data$Model))

ggplot(ci_data, aes(x = Model, y = Mean, color = Group)) +
  geom_point(size = 3, position = position_dodge(width = 0.5)) +  # Mean points
  geom_errorbar(aes(ymin = Lower_CI, ymax = Upper_CI), width = 0.2, position = position_dodge(width = 0.5)) +  # CI bars
  scale_color_manual(values = group_color_mapping) +  # Custom colors for groups
  facet_grid(Prediction ~ ., scales = "free_y") +  # Stack "Not Predicted" above "Predicted"
  coord_flip() +  # Flip axes for better readability
  labs(
    title = "Mean WTP Naturalness Estimates with 90% Confidence Intervals",
    x = "Model",
    y = "Mean Estimate (€/month)",
    color = "Group"
  ) +
  theme_minimal() +
  theme(
    legend.position = "right",
    strip.text.y = element_text(size = 12, face = "bold"),  # Adjust facet labels
    axis.text.y = element_text(size = 10)  # Ensure readable model names
  )