knitr::opts_chunk$set(
  echo    = FALSE,
  message = FALSE,
  warning = FALSE,
  fig.width  = 10,
  fig.height = 5,
  dpi     = 150
)
load("STAR_Students.RData")
library(dplyr)
library(ggplot2)
library(patchwork)
library(lme4)
library(emmeans)
library(tidyr)
library(forcats)
library(knitr)
library(scales)
library(data.table)
library(MatchIt)
library(purrr)
library(tibble)
library(nnet)
library(kableExtra)
library(plotly)

socio_econ_df = x %>% select(stdntid, race, gkclasstype, g1classtype, gksurban, g1surban, gkfreelunch, g1freelunch, gktreadss, g1treadss, gkschid, g1schid)

long_df <- socio_econ_df %>%
  pivot_longer(
    cols         = -c(stdntid, race),
    names_pattern = "(g[k1])(.+)",
    names_to     = c("grade", ".value")
  ) %>%
  mutate(
    classtype = fct_recode(as.factor(classtype),
      "Small"        = "SMALL CLASS",
      "Regular"      = "REGULAR CLASS",
      "Reg + Aide" = "REGULAR + AIDE CLASS"
    ),
    freelunch = as.factor(freelunch),
    race      = as.factor(race),
    surban    = as.factor(surban)
  )
long_df %>%
  group_by(classtype) %>%
  summarise(
    n_total   = n(),
    n_missing = sum(is.na(treadss)),
    pct_miss  = round(100 * mean(is.na(treadss)), 1)
  ) %>%
    kable(
    col.names = c("Class Type", "Total N", "Missing", "Percent Missing"),
    caption   = "Table 1: Missing Data in Reading Score by Class Type"
  )
long_df %>%
  filter(!is.na(classtype), !is.na(treadss)) %>%
  mutate(grade = fct_rev(factor(grade))) %>%
  ggplot(aes(x = treadss, fill = classtype)) +
  geom_density(alpha = 0.5) +
  facet_wrap(~grade) +
  labs(
    title    = "Visualization 1: Reading Score Distribution by Class Type",
    x = "Reading Score", y = "Density"
  ) +
  theme_bw()
long_df %>%
  filter(grade == "gk", !is.na(classtype)) %>%
  group_by(classtype) %>%
  summarise(
    n = n(),
    pct_black = round(100 * mean(race == "BLACK", na.rm = TRUE), 1),
    pct_freelunch = round(100 * mean(freelunch == "FREE LUNCH", na.rm = TRUE), 1),
    pct_urban = round(100 * mean(surban == "INNER CITY", na.rm = TRUE), 1),
    mean_read = round(mean(treadss, na.rm = TRUE), 1)
  ) %>%
    kable(
    col.names = c("Class Type", "Total", "Percent Black", "Percent Free Lunch", "Percent Inner City", "Mean Reading Score"),
    caption   = "Table 2: Distribution Check for variables for Kindergarten"
  )
long_df %>%
  filter(grade == "g1", !is.na(classtype)) %>%
  group_by(classtype) %>%
  summarise(
    n = n(),
    pct_black = round(100 * mean(race == "BLACK", na.rm = TRUE), 1),
    pct_freelunch = round(100 * mean(freelunch == "FREE LUNCH", na.rm = TRUE), 1),
    pct_urban = round(100 * mean(surban == "INNER CITY", na.rm = TRUE), 1),
    mean_read = round(mean(treadss, na.rm = TRUE), 1)
  ) %>%
    kable(
    col.names = c("Class Type", "Total", "Percent Black", "Percent Free Lunch", "Percent Inner City", "Mean Reading Score"),
    caption   = "Table 3: Distribution Check for variables for First Grade"
  )

long_df %>%
  filter(!is.na(classtype), !is.na(treadss)) %>%
  mutate(grade = fct_rev(factor(grade))) %>%
  ggplot(aes(x = classtype, y = treadss, fill = classtype)) +
  geom_boxplot(alpha = 0.6) +
  facet_wrap(~grade) +
  labs(
    title    = "Visualization 2: Reading Scores by Class Type",
    x = "Class Type", y = "Reading Score", fill = "Class Type"
  ) +
  theme_bw()

race_df <- x %>%
  filter(!is.na(race)) %>%
  count(race) %>%
  mutate(
    percent = n / sum(n),
    label = ifelse(percent > 0.02,
                   paste0(race, "\n", percent(percent, accuracy = 0.1)),
                   "")
  )


race_colors <- c(
  "WHITE" = "#4E79A7",
  "BLACK" = "#E15759",
  "ASIAN" = "#59A14F",
  "HISPANIC" = "#F28E2B",
  "OTHER" = "#B07AA1",
  "NATIVE AMERICAN" = "#76B7B2"
)

p_race <- ggplot(race_df, aes(x = 2, y = n, fill = race)) +
  geom_bar(stat = "identity", width = 1, color = "white") +
  coord_polar(theta = "y") +
  xlim(0.5, 2.5) +     
  geom_text(aes(label = label),
            position = position_stack(vjust = 0.5),
            size = 4) +
  scale_fill_manual(values = race_colors) +
  labs(
    title = "Donut Chart of Race Distribution",
    fill = "Race"
  ) +
  theme_void()

print(p_race)
plot_df <- x %>%
  filter(race %in% c("WHITE", "BLACK")) %>%
  pivot_longer(
    c(gktreadss, g1treadss),
    names_to = "grade",
    values_to = "reading_score"
  ) %>%
  mutate(
    class_label = case_when(
      grade == "gktreadss" ~ as.character(gkclasstype),
      grade == "g1treadss" ~ as.character(g1classtype)
    ),
    grade = factor(grade, levels = c("gktreadss", "g1treadss"),
                   labels = c("GK", "G1")),
    class_label = factor(
      class_label,
      levels = c("SMALL CLASS", "REGULAR CLASS", "REGULAR + AIDE CLASS"),
      labels = c("Small\nClass", "Regular\nClass", "Regular\n+ Aide")
    )
  ) %>%
  filter(!is.na(reading_score), !is.na(class_label))

ggplot(plot_df, aes(class_label, reading_score, fill = race)) +
  geom_boxplot(position = position_dodge(0.75), outlier.size = 0.8) +
  facet_wrap(~grade, nrow = 1) +
  scale_fill_manual(values = c("WHITE" = "#4E79A7", "BLACK" = "#E15759")) +
  labs(x = "Class Type", y = "Reading Score", fill = "Race") +
  theme_bw(base_size = 14) +
  theme(legend.position = "top")
long_df %>%
  filter(!is.na(freelunch), !is.na(treadss)) %>%
  mutate(grade = fct_rev(factor(grade))) %>%
  ggplot(aes(x = freelunch, y = treadss, fill = freelunch)) +
  geom_boxplot(alpha = 0.6) +
  facet_wrap(~grade) +
  labs(
    title    = "Visualization 4:Reading Scores by Free Lunch Status",
    x = "Free Lunch Status", y = "Reading Score"
  ) +
  theme_bw() 
long_df %>%
  filter(!is.na(surban), !is.na(treadss)) %>%
  mutate(grade = fct_rev(factor(grade))) %>%
  ggplot(aes(x = surban, y = treadss, fill = surban)) +
  geom_boxplot(alpha = 0.6) +
  facet_wrap(~grade) +
  labs(
    title = "Visualization 5:Reading Scores Across City Level",
    x = "City Level", y = "Reading Score"
  ) +
  theme_bw()
socio_econ_df = x %>% select(stdntid, gkclasstype, g1classtype, race, gksurban, g1surban, gkfreelunch, g1freelunch, gkschid, g1schid, gktreadss, g1treadss)

socio_econ_df = na.omit(socio_econ_df)
socio_econ_df[, 2:10] <- lapply(socio_econ_df[, 2:10], function(x) {
  factor(x, levels = unique(x), ordered = FALSE)
})
socio_econ_df = socio_econ_df %>% filter(race == "WHITE" | race == "BLACK")
socio_econ_df$g1classtype <- factor(socio_econ_df$g1classtype, 
                                     levels = c("SMALL CLASS",
                                                "REGULAR CLASS",
                                                "REGULAR + AIDE CLASS"))
ggplot(socio_econ_df, aes(x = reorder(gkschid, gktreadss, FUN = mean), 
                          y = gktreadss, 
                          fill = gksurban)) +
  geom_boxplot(width = 1) +  
  labs(x = "School", y = "Reading Score", title = "School vs Kindergarten Reading Scores") +
  theme_minimal() +
  theme(axis.text.x = element_blank())
ggplot(socio_econ_df, aes(x = reorder(g1schid, g1treadss, FUN = mean), 
                          y = g1treadss, 
                          fill = g1surban)) +
  geom_boxplot(width = 1) +  
  labs(x = "School", y = "Reading Score", title = "School vs First Grade Reading Scores")+
  theme_minimal() +
  theme(axis.text.x = element_blank())
my_table <- table(socio_econ_df$gkfreelunch, socio_econ_df$gksurban)
my_table_1 <- table(socio_econ_df$g1freelunch, socio_econ_df$g1surban)

chisq.test(my_table)
chisq.test(my_table_1)
data_gk_g1 <- x %>%
  select(gkclasstype, g1classtype) %>%
  na.omit()

flow_data <- data_gk_g1 %>%
  group_by(gkclasstype, g1classtype) %>%
  summarise(Freq = n(), .groups = "drop")

flow_data <- flow_data %>%
  mutate(
    gk_class = case_when(
      gkclasstype == "REGULAR CLASS" ~ "regular",
      gkclasstype == "SMALL CLASS" ~ "small",
      gkclasstype == "REGULAR + AIDE CLASS" ~ "regular+aide",
      TRUE ~ as.character(gkclasstype)
    ),
    g1_class = case_when(
      g1classtype == "REGULAR CLASS" ~ "regular",
      g1classtype == "SMALL CLASS" ~ "small",
      g1classtype == "REGULAR + AIDE CLASS" ~ "regular+aide",
      TRUE ~ as.character(g1classtype)
    )
  )

flow_data <- flow_data %>%
  mutate(
    source = paste0(gk_class, "_GK"),
    target = paste0(g1_class, "_G1")
  )

flow_data <- flow_data %>%
  mutate(
    color = case_when(
      gk_class == "regular" ~ "lightblue",
      gk_class == "small" ~ "salmon",
      gk_class == "regular+aide" ~ "lightgray",
      TRUE ~ "gray"
    )
  )

links <- flow_data %>%
  select(source, target, Freq, color)

node_order <- c(
  "regular_GK", "small_GK", "regular+aide_GK",
  "regular_G1", "small_G1", "regular+aide_G1"
)

nodes <- data.frame(name = node_order)


links$source_id <- match(links$source, nodes$name) - 1
links$target_id <- match(links$target, nodes$name) - 1

node_labels <- c(
  "regular GK", "small GK", "reg+aide GK",
  "regular G1", "small G1", "reg+aide G1"
)


node_colors <- c(
  "#EFC000FF", "#868686FF", "#CD534CFF",
  "#EFC000FF", "#868686FF", "#CD534CFF"
)

fig <- plot_ly(
  type = "sankey",
  orientation = "h",
  node = list(
    label = node_labels,
    color = node_colors,
    pad = 20,
    thickness = 20,
    line = list(color = "black", width = 0.5)
  ),
  link = list(
    source = links$source_id,
    target = links$target_id,
    value  = links$Freq,
    color  = links$color
  )
)

fig <- fig %>%
  layout(
    title = "Continuity of Program from Kindergarten to Grade 1",
    font = list(size = 14),
    margin = list(t = 80),
    width = 600
  )

fig
model_k <- lmer(gktreadss ~ gkclasstype + gkfreelunch + race + 
                          (1 | gkschid), data = socio_econ_df, REML = FALSE)
anova(model_k)
model_1 <- lmer(g1treadss ~ g1classtype + g1freelunch + race +
                          (1 | g1schid), data = socio_econ_df, REML = FALSE)
anova(model_1)
group_means_k = emmeans(model_k, "gkclasstype", pbkrtest.limit = 4000)
pairs(group_means_k, adjust = "tukey")
group_means_1 = emmeans(model_1, "g1classtype", pbkrtest.limit = 4000)
pairs(group_means_1, adjust = "tukey")
socio_econ_df$gkclasstype <- relevel(socio_econ_df$gkclasstype, ref = "SMALL CLASS")
model_k <- lmer(gktreadss ~ gkclasstype + gkfreelunch + race + (1|gkschid), 
                data = socio_econ_df)

print(VarCorr(model_k))
summary(model_k)$coefficients
socio_econ_df$g1classtype <- relevel(socio_econ_df$g1classtype, ref = "SMALL CLASS")
print(VarCorr(model_1))

summary(model_1)$coefficients
res_df <- bind_rows(
  data.frame(
    fitted = fitted(model_k),
    resid  = resid(model_k),
    grade  = "GK"
  ),
  data.frame(
    fitted = fitted(model_1),
    resid  = resid(model_1),
    grade  = "G1"
  )
) %>%
  mutate(grade = factor(grade, levels = c("GK","G1"))) 

ggplot(res_df, aes(x = fitted, y = resid)) +
  geom_point(alpha = 0.35, size = 1.5) +
  geom_hline(yintercept = 0, linetype = "dashed", color = "gray40") +
  geom_smooth(method = "loess", se = FALSE, color = "red", linewidth = 1) +
  facet_wrap(~grade, nrow = 1, scales = "free") +  
  labs(
    x = "Fitted Values",
    y = "Residuals",
    title = "Residual Plots for GK and G1 Mixed Models"
  ) +
  theme_bw(base_size = 14) +
  theme(
    strip.text = element_text(size = 14),
    plot.title = element_text(hjust = 0.5)
  )

test_independence <- function(data, class, school, covar, keep = NULL) {
  
  df <- data %>%
    select(all_of(c(class, school, covar))) %>%
    na.omit()
  
  if (!is.null(keep)) {
    df <- df %>% filter(.data[[covar]] %in% keep)
  }
  
  df <- df %>% mutate(across(everything(), as.factor))
  
  m1 <- multinom(as.formula(paste(class, "~", school)), data=df, trace=FALSE)
  m2 <- multinom(as.formula(paste(class, "~", school, "+", covar)), data=df, trace=FALSE)
  
  LR <- 2*(logLik(m2) - logLik(m1))
  df_diff <- attr(logLik(m2),"df") - attr(logLik(m1),"df")
  p <- pchisq(LR, df_diff, lower.tail = FALSE)
  
  c(LR = as.numeric(LR), df = df_diff, p = p)
}

results <- rbind(
  "GK: Class Type vs Race" =
    test_independence(x, "gkclasstype","gkschid","race", keep=c("WHITE","BLACK")),
  
  "GK: Class Type vs Free Lunch" =
    test_independence(x,"gkclasstype","gkschid","gkfreelunch"),
  
  "G1: Class Type vs Race" =
    test_independence(x,"g1classtype","g1schid","race", keep=c("WHITE","BLACK")),
  
  "G1: Class Type vs Free Lunch" =
    test_independence(x,"g1classtype","g1schid","g1freelunch")
)

results <- as.data.frame(results)

knitr::kable(
  results,
  digits = 4,
  caption = "Likelihood ratio tests for independence between class assignment and covariates within schools"
)
library(dplyr)
library(purrr)
library(tidyr)
library(MatchIt)
library(ggplot2)

df <- x %>%
  select(stdntid, race, gkschid, gkclasstype, g1classtype,
         gksurban, gkfreelunch, gktreadss, g1treadss) %>%
  drop_na() %>%
  mutate(
    pctR_gk = 100 * percent_rank(gktreadss),
    delta_read = g1treadss - gktreadss
  )

run_matching <- function(data, origin, dest) {
  temp <- data %>%
    mutate(
      treat = case_when(
        gkclasstype == origin & g1classtype == dest   ~ 1,
        gkclasstype == origin & g1classtype == origin ~ 0,
        TRUE ~ NA_real_
      )
    ) %>%
    filter(!is.na(treat))

  if (nrow(temp) < 20 || sum(temp$treat == 1) < 10 || sum(temp$treat == 0) < 10) return(NULL)

  m <- matchit(
    treat ~ pctR_gk + race + gkfreelunch,
    data = temp,
    method = "nearest",
    distance = "glm",
    ratio = 1,
    caliper = 0.2,
    std.caliper = TRUE,
    exact = ~ gkschid
  )

  matched <- match.data(m)

  means <- matched %>%
    group_by(treat) %>%
    summarise(mean_change = weighted.mean(delta_read, weights), .groups = "drop")

  tt <- tryCatch(t.test(delta_read ~ treat, data = matched), error = function(e) NULL)

  tibble(
    origin = origin,
    dest = dest,
    ATT = means$mean_change[means$treat == 1] - means$mean_change[means$treat == 0],
    treated_mean = means$mean_change[means$treat == 1],
    control_mean = means$mean_change[means$treat == 0],
    treated_n = sum(matched$treat == 1),
    control_n = sum(matched$treat == 0),
    p_value = ifelse(is.null(tt), NA, tt$p.value)
  )
}

transitions <- tribble(
  ~origin, ~dest,
  "SMALL CLASS", "REGULAR CLASS",
  "SMALL CLASS", "REGULAR + AIDE CLASS",
  "REGULAR CLASS", "SMALL CLASS",
  "REGULAR CLASS", "REGULAR + AIDE CLASS",
  "REGULAR + AIDE CLASS", "SMALL CLASS",
  "REGULAR + AIDE CLASS", "REGULAR CLASS"
)

results <- pmap_dfr(transitions, ~run_matching(df, ..1, ..2)) %>%
  mutate(
    transition = paste(origin, "→", dest),
    sig = ifelse(!is.na(p_value) & p_value < 0.05, "Significant", "Not significant")
  )

p_readchange <- ggplot(results, aes(x = reorder(transition, ATT), y = ATT, fill = sig)) +
  geom_col() +
  coord_flip() +
  geom_hline(yintercept = 0, linetype = "dashed") +
  theme_minimal() +
  labs(
    x = "Class transition",
    y = "ATT (G1 reading - GK reading)",
    title = "Estimated effect of class transitions on reading score change",
    fill = NULL
  ) +
  theme(
    axis.text.y = element_text(color = "black", size = 10, face = "bold"),
    plot.title = element_text(hjust = 0.5)
  )

print(p_readchange)
change_df <- x %>%
  select(stdntid,
         gkclasstype, g1classtype,
         race,
         gkfreelunch, g1freelunch,
         gkschid, g1schid,
         gktreadss, g1treadss) %>%
  na.omit()

change_df <- change_df %>%
  filter(race %in% c("WHITE", "BLACK"))

change_df <- change_df %>%
  mutate(delta_reading = g1treadss - gktreadss)

# factor coding
change_df$gkclasstype <- factor(
  change_df$gkclasstype,
  levels = c("SMALL CLASS", "REGULAR CLASS", "REGULAR + AIDE CLASS")
)

change_df$race <- factor(change_df$race)
change_df$gkfreelunch <- factor(change_df$gkfreelunch)
change_df$gkschid <- factor(change_df$gkschid)
change_df <- x %>%
  select(
    stdntid,
    gkclasstype, g1classtype,
    race,
    gkfreelunch, g1freelunch,
    gkschid, g1schid,
    gktreadss, g1treadss
  ) %>%
  na.omit() %>%
  filter(race %in% c("WHITE", "BLACK")) %>%
  mutate(
    delta_reading = g1treadss - gktreadss,
    gkclasstype = factor(
      gkclasstype,
      levels = c("SMALL CLASS", "REGULAR CLASS", "REGULAR + AIDE CLASS")
    ),
    gkfreelunch = factor(
      gkfreelunch,
      levels = c("FREE LUNCH", "NON-FREE LUNCH")
    ),
    race = factor(
      race,
      levels = c("BLACK", "WHITE")
    ),
    gkschid = factor(gkschid)
  )

change_model <- lmer(
  delta_reading ~ gkclasstype + gkfreelunch + race + (1 | gkschid),
  data = change_df,
  REML = FALSE
)

anova(change_model)