library(ggplot2)
library(dplyr)
library(rstanarm)  # For Bayesian modeling
library(bayesplot) # For posterior visualization
library(tidybayes) # For summarizing posterior distributions
library(TOSTER)    # For equivalence testing
library(brms)      # For more complex Bayesian GLMs
library(broom)     # For cleaning model output for plots

# Parallel processing
options(mc.cores = parallel::detectCores())

# Load data
data <- read.csv("/Users/devynmiller/Downloads/hw2.csv")

theme_set(theme_minimal())
# Figure 1: Distribution of Anxiety Scores by Group
ggplot(data, aes(x = group, y = anxiety, fill = group)) +
  geom_violin(alpha = 0.6) +
  geom_boxplot(width = 0.2, alpha = 0.8) +
  labs(
    title = "Figure 1: Distribution of Anxiety Scores by Group",
    subtitle = "Violin plot with overlaid boxplot showing group-level anxiety distributions",
    x = "Group (A vs. B)",
    y = "Anxiety Score"
  ) +
  scale_fill_manual(values = c("A" = "red", "B" = "blue"), labels = c("Group A", "Group B")) +
  theme(legend.position = "none")

# Figure 2: Frequentist Confidence Interval for Group Difference in Anxiety
group_A <- data %>% filter(group == "A") %>% pull(anxiety)
group_B <- data %>% filter(group == "B") %>% pull(anxiety)
mean_diff <- mean(group_A) - mean(group_B)
t_test_result <- t.test(group_A, group_B, var.equal = TRUE)
ci_lower <- t_test_result$conf.int[1]
ci_upper <- t_test_result$conf.int[2]

# Simulate distribution for Frequentist confidence interval
freq_dist <- rnorm(10000, mean = mean_diff, sd = sd(c(group_A, group_B)) / sqrt(length(group_A) + length(group_B) - 2))
freq_density <- density(freq_dist)
freq_density_df <- data.frame(x = freq_density$x, y = freq_density$y)

ggplot(freq_density_df, aes(x = x, y = y)) +
  geom_line(color = "blue") +
  geom_area(data = subset(freq_density_df, x >= ci_lower & x <= ci_upper), fill = "lightblue", alpha = 0.5) +
  geom_vline(xintercept = mean_diff, linetype = "dashed", color = "red") +
  annotate("text", x = mean_diff, y = max(freq_density$y) * 0.9, label = paste("Mean =", round(mean_diff, 2)), color = "red", hjust = -0.1, vjust = -1.5) +
  geom_vline(xintercept = ci_lower, linetype = "dashed", color = "black") +
  geom_vline(xintercept = ci_upper, linetype = "dashed", color = "black") +
  annotate("text", x = ci_lower, y = max(freq_density$y) * 0.9, label = paste("Lower 95% Confidence Interval =", round(ci_lower, 2)), color = "black", hjust = .2, vjust = -1.5) +
  annotate("text", x = ci_upper, y = max(freq_density$y) * 0.9, label = paste("Upper 95% Confidence Interval =", round(ci_upper, 2)), color = "black", hjust = .8, vjust = -1.5) +
  labs(
    title = "Figure 2: Frequentist Confidence Interval for Group Difference in Anxiety",
    subtitle = "Distribution of mean difference with 95% confidence interval",
    x = "Mean Difference in Anxiety (A - B)",
    y = "Density"
  )

cat("Figure 2 Summary: Mean difference in anxiety between groups A and B =", round(mean_diff, 2), 
    "with 95% CI [", round(ci_lower, 2), ",", round(ci_upper, 2), "]\n")

# Figure 3: Bayesian Posterior Distribution for Group Difference in Anxiety

bayesian_model_q1 <- brm(
  anxiety ~ group,
  data = data,
  family = gaussian(),
  prior = c(prior(normal(0, 5), class = "Intercept"),
            prior(normal(0, 2.5), class = "b")),
  chains = 4,
  iter = 4000
)


# Convergence Diagnostics
mcmc_trace(bayesian_model_q1)

# Posterior Predictive Check
pp_check(bayesian_model_q1)


# Summary and Posterior Distribution Plot
posterior_samples_q1 <- posterior_samples(bayesian_model_q1, "b_groupB")
posterior_density <- density(posterior_samples_q1$b_groupB)
posterior_density_df <- data.frame(x = posterior_density$x, y = posterior_density$y)
credible_interval_lower <- quantile(posterior_samples_q1$b_groupB, 0.025)
credible_interval_upper <- quantile(posterior_samples_q1$b_groupB, 0.975)

ggplot(posterior_density_df, aes(x = x, y = y)) +
  geom_line(color = "darkgreen") +
  geom_area(data = subset(posterior_density_df, x >= credible_interval_lower & x <= credible_interval_upper), fill = "lightgreen", alpha = 0.5) +
  geom_vline(xintercept = 0, linetype = "dashed", color = "red") +  # Null effect line
  geom_vline(xintercept = mean(posterior_samples_q1$b_groupB), linetype = "dashed", color = "blue") +  # Observed mean effect
  annotate("text", x = 0, y = max(posterior_density$y) * 0.9, label = "Null Effect", color = "red", hjust = 1) +
  annotate("text", x = mean(posterior_samples_q1$b_groupB), y = max(posterior_density$y) * 0.9, label = paste("Mean =", round(mean(posterior_samples_q1$b_groupB), 2)), color = "blue", hjust = -0.1) +
  geom_vline(xintercept = credible_interval_lower, linetype = "dashed", color = "black") +
  geom_vline(xintercept = credible_interval_upper, linetype = "dashed", color = "black") +
  annotate("text", x = credible_interval_lower, y = max(posterior_density$y) * 0.9, label = paste("Lower 95% Credible Interval =", round(credible_interval_lower, 2)), color = "black", hjust = .5, vjust = -1.5) +
  annotate("text", x = credible_interval_upper, y = max(posterior_density$y) * 0.9, label = paste("Upper 95% Credible Interval =", round(credible_interval_upper, 2)), color = "black", hjust = .5, vjust = -1.5) +
  labs(
    title = "Figure 3: Bayesian Posterior Distribution of Group Difference in Anxiety",
    subtitle = "95% credible interval for group difference",
    x = "Difference in Anxiety (Group B - Group A)",
    y = "Density"
  )
cat("Figure 3 Summary: Bayesian posterior mean difference =", round(mean(posterior_samples_q1$b_groupB), 2), 
    "with 95% credible interval [", round(credible_interval_lower, 2), ",", round(credible_interval_upper, 2), "]\n")


# Figure 6: Frequentist Distribution Plot for Effect of Corgi Petting
glm_model_q2 <- lm(mental_well ~ corgi + married + sex + anxiety + age, data = data)
corgi_effect <- coef(summary(glm_model_q2))["corgi", "Estimate"]
corgi_se <- coef(summary(glm_model_q2))["corgi", "Std. Error"]
corgi_ci_lower <- corgi_effect - 1.96 * corgi_se
corgi_ci_upper <- corgi_effect + 1.96 * corgi_se

freq_dist_corgi <- rnorm(10000, mean = corgi_effect, sd = corgi_se)
freq_density_corgi <- density(freq_dist_corgi)
freq_density_corgi_df <- data.frame(x = freq_density_corgi$x, y = freq_density_corgi$y)

ggplot(freq_density_corgi_df, aes(x = x, y = y)) +
  geom_line(color = "blue") +
  geom_area(data = subset(freq_density_corgi_df, x >= corgi_ci_lower & x <= corgi_ci_upper), fill = "lightblue", alpha = 0.5) +
  geom_vline(xintercept = corgi_effect, linetype = "dashed", color = "red") +
  annotate("text", x = corgi_effect, y = max(freq_density_corgi$y) * 0.9, label = paste("Mean =", round(corgi_effect, 2)), color = "red", hjust = -0.1) +
  geom_vline(xintercept = corgi_ci_lower, linetype = "dashed", color = "black") +
  geom_vline(xintercept = corgi_ci_upper, linetype = "dashed", color = "black") +
  annotate("text", x = corgi_ci_lower, y = max(freq_density_corgi$y) * 0.9, label = paste("Lower 95% Confidence Interval =", round(corgi_ci_lower, 2)), color = "black", hjust = .5, vjust = -1.5) +
  annotate("text", x = corgi_ci_upper, y = max(freq_density_corgi$y) * 0.9, label = paste("Upper 95% Confidence Interval =", round(corgi_ci_upper, 2)), color = "black", hjust = .5, vjust = -1.5) +
  labs(
    title = "Figure 6: Frequentist Distribution for Effect of Corgi Petting",
    subtitle = "Distribution of estimated effect with 95% confidence interval",
    x = "Effect of Corgi Petting on Mental Well-being",
    y = "Density"
  )

# Frequentist Approach: Corgi effect point estimate and 95% CI
corgi_effect <- coef(summary(glm_model_q2))["corgi", "Estimate"]
corgi_se <- coef(summary(glm_model_q2))["corgi", "Std. Error"]
corgi_ci_lower <- corgi_effect - 1.96 * corgi_se
corgi_ci_upper <- corgi_effect + 1.96 * corgi_se
cat("Frequentist estimate of corgi effect:", round(corgi_effect, 2), 
    "with 95% CI [", round(corgi_ci_lower, 2), ",", round(corgi_ci_upper, 2), "]\n")



# Figure 7: Bayesian Posterior Distribution of Corgi Effect on Mental Well-being
bayesian_glm_q2 <- brm(
  mental_well ~ corgi + married + sex + anxiety + age,
  data = data,
  family = gaussian(),
  prior = prior(normal(0, 5), class = "b"),
  chains = 4,
  iter = 4000
)

# Convergence Diagnostics
mcmc_trace(bayesian_glm_q2)

# Posterior Predictive Check
pp_check(bayesian_glm_q2)

posterior_samples_q2 <- posterior_samples(bayesian_glm_q2, "b_corgi")
posterior_density_q2 <- density(posterior_samples_q2$b_corgi)
posterior_density_q2_df <- data.frame(x = posterior_density_q2$x, y = posterior_density_q2$y)
credible_interval_lower_q2 <- quantile(posterior_samples_q2$b_corgi, 0.025)
credible_interval_upper_q2 <- quantile(posterior_samples_q2$b_corgi, 0.975)

ggplot(posterior_density_q2_df, aes(x = x, y = y)) +
  geom_line(color = "darkgreen") +
  geom_area(data = subset(posterior_density_q2_df, x >= credible_interval_lower_q2 & x <= credible_interval_upper_q2), fill = "lightgreen", alpha = 0.5) +
  geom_vline(xintercept = mean(posterior_samples_q2$b_corgi), linetype = "dashed", color = "blue") +  # Observed mean effect
  annotate("text", x = 0, y = max(posterior_density_q2$y) * 0.9, label = " ", color = "red", hjust = -0.1) +
  annotate("text", x = mean(posterior_samples_q2$b_corgi), y = max(posterior_density_q2$y) * 0.9, label = paste("Mean =", round(mean(posterior_samples_q2$b_corgi), 2)), color = "blue", hjust = -0.1) +
  geom_vline(xintercept = credible_interval_lower_q2, linetype = "dashed", color = "black") +
  geom_vline(xintercept = credible_interval_upper_q2, linetype = "dashed", color = "black") +
  annotate("text", x = credible_interval_lower_q2, y = max(posterior_density_q2$y) * 0.9, label = paste("Lower 95% Credible Interval =", round(credible_interval_lower_q2, 2)), color = "black", hjust =.5, vjust = -1.5) +
  annotate("text", x = credible_interval_upper_q2, y = max(posterior_density_q2$y) * 0.9, label = paste("Upper 95% Credible Interval =", round(credible_interval_upper_q2, 2)), color = "black", hjust = .5, vjust = -1.5) +
  labs(
    title = "Figure 7: Bayesian Posterior Distribution of Corgi Effect on Mental Well-being",
    subtitle = "95% credible interval for the effect of petting corgis on well-being",
    x = "Effect of Corgi Petting on Mental Well-being",
    y = "Density"
  )
# Bayesian Approach: Corgi effect point estimate and 95% Credible Interval
posterior_samples_q2 <- posterior_samples(bayesian_glm_q2, "b_corgi")
credible_interval_lower_q2 <- quantile(posterior_samples_q2$b_corgi, 0.025)
credible_interval_upper_q2 <- quantile(posterior_samples_q2$b_corgi, 0.975)
cat("Bayesian posterior estimate of corgi effect:", round(mean(posterior_samples_q2$b_corgi), 2), 
    "with 95% credible interval [", round(credible_interval_lower_q2, 2), ",", round(credible_interval_upper_q2, 2), "]\n")