library(MASS)
library(ggplot2)
library(lme4)
# library(lmerTest) # used for lecture 12
library(sjPlot)


set.seed(123)

# -----------------------------
# 1) Design: 5 programs, 10–20 students each
# -----------------------------
m <- 5
n_m <- sample(10:20, size = m, replace = TRUE)
program <- factor(rep(paste0("Prog", 1:m), times = n_m))
n <- length(program)

# x independent of program
x <- pmax(0, rnorm(n, mean = 6, sd = 2))  # study hours, truncated at 0
dat <- data.frame(program = program, x = x)

# -----------------------------
# 2) True fixed effects (population-average relationship)
# -----------------------------
beta0 <- 40     # baseline test score
beta1 <- 5.0    # increase per study hour

# -----------------------------
# 3) Random effects: (g0_i, g1_i) ~ MVN(0, Q)
# -----------------------------
sd_g0 <- 6.0        # SD of random intercepts
sd_g1 <- 1.2        # SD of random slopes

cor_g0g1 <- -0.3    # correlation between intercept and slope

Q <- matrix(c(sd_g0^2,               cor_g0g1*sd_g0*sd_g1,
              cor_g0g1*sd_g0*sd_g1,  sd_g1^2),
            nrow = 2, byrow = TRUE)

# Residual SD 
sigma_eps <- 8.0

# Draw random effects per program
gammavec <- mvrnorm(n = m, mu = c(0, 0), Sigma = Q)
colnames(gammavec) <- c("g0", "g1")
re <- data.frame(program = factor(paste0("Prog", 1:m), levels = levels(program)),
                 g0 = gammavec[, "g0"], g1 = gammavec[, "g1"])

# Merge RE into observation-level data
dat <- merge(dat, re, by = "program", sort = FALSE)

# -----------------------------
# 4) Generate outcome
# -----------------------------
dat$y <- (beta0 + dat$g0) + (beta1 + dat$g1) * dat$x + rnorm(n, 0, sigma_eps)

# Keep only analysis columns (drop the true RE columns)
dat_sim <- dat[, c("y", "x", "program")]

# 1) Plot of x vs y (all points)
p1 <- ggplot(dat_sim, aes(x = x, y = y)) +
  geom_point() +
  labs(x = "Study hours (x)", y = "Test score (y)",
       title = "x vs y (all students)")

# 2) Plot of x vs y, groups coloured
p2 <- ggplot(dat_sim, aes(x = x, y = y, color = program)) +
  geom_point() +
  labs(x = "Study hours (x)", y = "Test score (y)",
       title = "x vs y by study program", color = "Program")

print(p1)
print(p2)

p_faceted <- ggplot(dat_sim, aes(x = x, y = y)) +
  geom_point() +
  geom_smooth(method = "lm", se = FALSE) +   # simple linear slope per facet
  facet_wrap(~ program) +
  labs(x = "Study hours (x)",
       y = "Test score (y)",
       title = "Study hours vs Test score by program\n(Simple linear regression per program)")

print(p_faceted)

# Try dummy coding for different intercepts and means
dat_sim$program <- relevel(dat_sim$program, ref = "Prog1")
lm_full <- lm(y ~ x * program, data = dat_sim)
summary(lm_full)


# -----------------------------
# 5) Fit LMM: random intercept + random slope
# -----------------------------

fit <- lmer(y ~ x + (1 + x | program), data = dat_sim, REML = TRUE)

summary(fit)

# Random effects (predicted) per program
ranef(fit)$program

# Fixed effects (population level)
fixef(fit)

plot_model(fit, type = "re")