data.frame(rho = c(0, 0.3, 0.7)) |>
  dplyr::mutate(Var_Delta = 2*(1-rho),
                SE_Delta  = sqrt(Var_Delta))

set.seed(667)
library(dplyr)
library(tidyr)
library(ggplot2)
library(nlme)

# Simulate balanced longitudinal data with AR(1) correlation + heteroscedasticity
N   <- 120           # subjects
m   <- 5             # visits
times <- 0:4
rho_true <- 0.6      # true serial correlation
sigma_true <- 1.2    # marginal SD scale for the AR(1) part
omega <- c(1.0, 1.1, 1.5, 1.2, 0.9)  # time-specific SD multipliers

# AR(1) generator with stationary Var = 1
sim_ar1_std <- function(m, rho){
  z <- numeric(m)
  z[1] <- rnorm(1, sd = 1)
  if (m > 1) {
    eps <- rnorm(m - 1, sd = sqrt(1 - rho^2))
    for (t in 2:m) z[t] <- rho * z[t - 1] + eps[t - 1]
  }
  z
}

ids <- factor(seq_len(N))
design <- expand.grid(id = ids, time_idx = seq_len(m)) |>
  arrange(id, time_idx) |>
  mutate(time_f   = factor(time_idx, levels = seq_len(m),
                           labels = paste0("t", times)),
         time_num = times[time_idx])

# Mean structure: linear time trend + small group effect
beta0 <- 10
beta_time <- 0.8
group <- factor(rep(c("A","B"), each = N/2))
design <- design |>
  group_by(id) |>
  mutate(group = group[as.integer(id)]) |>
  ungroup()

# Subject random intercept
sigma_b <- 1.0
b_i <- rnorm(N, 0, sigma_b)
b <- b_i[as.integer(design$id)]

# AR(1) errors with time-specific SD multipliers
z <- do.call(c, lapply(seq_len(N), function(i) sim_ar1_std(m, rho_true)))
err <- sigma_true * (omega[design$time_idx]) * z

sim <- design |>
  mutate(mu = beta0 + beta_time * time_num + ifelse(group=="B", 0.5, 0),
         y  = mu + b + err)

# Fit GLS under different covariance structures
fit_ind <- gls(y ~ time_num, data = sim, method = "REML")

fit_ar1 <- gls(y ~ time_num, data = sim,
               correlation = corAR1(form = ~ time_idx | id),
               method = "REML")

# Initialize varIdent with numeric starting values
w_time <- varIdent(form = ~ 1 | time_f)
nlme::coef(w_time) <- setNames(rep(1, nlevels(sim$time_f) - 1),
                               levels(sim$time_f)[-1])

fit_vhet <- gls(y ~ time_num, data = sim,
                weights = w_time,
                method  = "REML")

fit_arh1 <- gls(y ~ time_num, data = sim,
                correlation = corAR1(form = ~ time_idx | id),
                weights     = w_time,
                method      = "REML")

# Compare fits.
# REML-LRT (nested) and AIC (non-nested) are the primary criteria;
# BIC is shown for reference only (FLW p. 179 advises against it for Sigma).
tab_ic <- data.frame(
  model  = c("Independence","AR(1) only","varIdent only","ARH(1) = AR(1)+varIdent"),
  AIC    = c(AIC(fit_ind), AIC(fit_ar1), AIC(fit_vhet), AIC(fit_arh1)),
  BIC    = c(BIC(fit_ind), BIC(fit_ar1), BIC(fit_vhet), BIC(fit_arh1)),
  logLik = c(logLik(fit_ind), logLik(fit_ar1), logLik(fit_vhet), logLik(fit_arh1))
) |>
  arrange(AIC)

tab_ic

# Estimated AR(1) correlation
rho_hat_ar1  <- as.numeric(coef(fit_ar1$modelStruct$corStruct))
rho_hat_arh1 <- as.numeric(coef(fit_arh1$modelStruct$corStruct))

data.frame(
  param = c("rho_hat (AR1 only)","rho_hat (ARH1)"),
  estimate = c(rho_hat_ar1, rho_hat_arh1)
)

ggplot(sd_table, aes(x = time, y = sd_est, group = 1)) +
  geom_line() + geom_point(size = 2) +
  labs(title = "Estimated residual SD by time (GLS with AR(1) + varIdent)",
       x = "Time (factor levels)", y = "Residual SD") +
  theme_minimal()

library(dplyr)
library(tidyr)
library(nlme)
library(emmeans)
library(ggplot2)

# Read TLC data
path <- "../../data/tlc.csv"
tlc_wide <- read.csv(path, check.names = FALSE, stringsAsFactors = FALSE)
names(tlc_wide) <- trimws(names(tlc_wide))

tlc_wide <- tlc_wide %>%
  rename(
    id  = ID,
    trt = `Treatment Group`,
    w0  = `Lead Level Week 0`,
    w1  = `Lead Level Week 1`,
    w4  = `Lead Level Week 4`,
    w6  = `Lead Level Week 6`
  ) %>%
  mutate(trt = factor(trt, levels = c("P","A"), labels = c("Placebo","Succimer")))

# Wide -> long format
long <- tlc_wide %>%
  pivot_longer(c(w0,w1,w4,w6), names_to = "wk", values_to = "y") %>%
  mutate(
    time       = recode(wk, w0="0", w1="1", w4="4", w6="6"),
    time_f     = factor(time, levels = c("0","1","4","6")),
    time_num   = as.numeric(as.character(time)),
    time_idx   = as.integer(time_f),
    group      = trt,
    id         = factor(id),
    y          = as.numeric(y)
  ) %>%
  select(id, group, time_f, time_num, time_idx, y) %>%
  arrange(id, time_idx)

head(long)

# Helpers for stable varIdent initialization
make_varIdent <- function(data, fac) {
  f <- nlme::varIdent(form = as.formula(paste0("~ 1 | ", fac)))
  levs <- levels(data[[fac]])
  if (length(levs) > 1) nlme::coef(f) <- stats::setNames(rep(1, length(levs) - 1), levs[-1])
  f
}
w_time <- make_varIdent(long, "time_f")

m_visits <- nlevels(long$time_f)
p_lower  <- m_visits * (m_visits - 1) / 2
start_vec_un <- rep(0.2, p_lower)

# Check planned time spacing
planned <- levels(long$time_f)
gaps <- diff(sort(unique(long$time_num)))
data.frame(planned_times = planned, unique_gaps = paste(unique(gaps), collapse = ", "))

# Compute rho at baseline vs week 1
bl_w1 <- tlc_wide %>% select(id, trt, w0, w1) %>% filter(!is.na(w0) & !is.na(w1))
rho_hat <- with(bl_w1, cor(w0, w1))

var_within_hat   <- var(bl_w1$w1 - bl_w1$w0)
var_between_hat  <- var(bl_w1$w1) + var(bl_w1$w0)
ratio_hat <- var_within_hat / var_between_hat

data.frame(
  rho_hat = round(rho_hat, 3),
  ratio_within_to_between = round(ratio_hat, 3),
  approx_1_minus_rho = round(1 - rho_hat, 3)
)

set.seed(667)
library(dplyr)
library(tidyr)
library(ggplot2)
library(nlme)

# Simulate longitudinal data with AR(1) errors
N   <- 400      # more subjects
m   <- 8        # more visits
rho <- 0.80     # stronger serial correlation
sigma <- 2      # marginal SD

# AR(1) generator with stationary Var = sigma^2
sim_ar1 <- function(m, rho, sigma){
  e <- numeric(m)
  e[1] <- rnorm(1, sd = sigma)
  if (m > 1) {
    w <- rnorm(m - 1, sd = sigma * sqrt(1 - rho^2))
    for (t in 2:m) e[t] <- rho * e[t - 1] + w[t - 1]
  }
  e
}

ids   <- factor(seq_len(N))
times <- seq_len(m)

design <- expand.grid(id = ids, time_idx = times) |>
  arrange(id, time_idx) |>
  mutate(
    group    = ifelse(as.integer(id) <= N/2, "A", "B"),
    group    = factor(group, levels = c("A","B")),
    time_f   = factor(time_idx, levels = times),
    time_num = as.numeric(time_idx)
  )

# Fixed mean structure
beta0 <- 10
beta_time <- 0.7
beta_groupB <- 1.0
beta_int_B <- 0.3

errs <- do.call(c, lapply(seq_len(N), function(i) sim_ar1(m, rho, sigma)))

sim_long <- design |>
  mutate(
    mu = beta0 +
         beta_time   * (time_num - 1) +
         beta_groupB * as.numeric(group == "B") +
         beta_int_B  * as.numeric(group == "B") * (time_num - 1),
    y  = mu + errs
  )

# Fit GLS under Independence and AR(1)
gls_ind <- gls(y ~ group * time_f, data = sim_long, method = "REML")

gls_ar1 <- gls(y ~ group * time_f, data = sim_long,
               correlation = corAR1(form = ~ time_idx | id),
               method = "REML")

# Attach normalized residuals
res_ind <- sim_long |> mutate(res = residuals(gls_ind, type = "normalized"))
res_ar1 <- sim_long |> mutate(res = residuals(gls_ar1, type = "normalized"))

# Pooled residual ACF by subject (lags 1-3)
lag_cor <- function(x, k){
  n <- length(x); if (n <= k) return(NA_real_)
  stats::cor(x[(k+1):n], x[1:(n-k)], use = "pairwise.complete.obs")
}

acf_by_id <- function(dat){
  dat |>
    arrange(id, time_idx) |>
    group_by(id) |>
    summarise(acf1 = lag_cor(res, 1),
              acf2 = lag_cor(res, 2),
              acf3 = lag_cor(res, 3),
              .groups = "drop")
}

pool_acf <- function(dat, model_lab){
  by_id <- acf_by_id(dat)
  by_id |>
    tidyr::pivot_longer(cols = dplyr::starts_with("acf"),
                        names_to = "lag", values_to = "rho") |>
    mutate(lag = as.integer(sub("acf", "", lag))) |>
    group_by(lag) |>
    summarise(
      mean_rho = mean(rho, na.rm = TRUE),
      med_rho  = median(rho, na.rm = TRUE),
      se_rho   = sd(rho, na.rm = TRUE)/sqrt(dplyr::n()),
      .groups  = "drop"
    ) |>
    mutate(model = model_lab)
}

acf_ind <- pool_acf(res_ind, "Independence")
acf_ar1 <- pool_acf(res_ar1, "AR(1)")
acf_all <- bind_rows(acf_ind, acf_ar1)

ggplot(acf_all, aes(x = factor(lag), y = mean_rho, fill = model)) +
  geom_col(position = position_dodge(width = 0.7), width = 0.6) +
  geom_errorbar(aes(ymin = mean_rho - 2*se_rho, ymax = mean_rho + 2*se_rho),
                position = position_dodge(width = 0.7), width = 0.15) +
  geom_hline(yintercept = 0, linewidth = 0.3) +
  labs(title = "Pooled residual ACF by model (simulated AR(1) truth, rho = 0.80)",
       subtitle = "Correct Sigma (AR(1)) shows flat residual ACF; Independence shows positive lag-1 correlation",
       x = "Lag", y = "Residual correlation") +
  theme_minimal()

library(nlme)

# Calculate empirical variogram from independence model residuals
vario_gls <- Variogram(gls_ind, form = ~ time_num)

plot(vario_gls, main = "Empirical Semivariogram of Residuals",
     xlab = "Time Lag (Weeks)", ylab = "Semivariance")

# Fit the unstructured correlation model on the simulated AR(1) data
gls_un_sim <- gls(y ~ group * time_f, data = sim_long,
                  correlation = corSymm(form = ~ time_idx | id),
                  weights = varIdent(form = ~ 1 | time_f),
                  method = "REML")

# Estimated UN correlation matrix (the "data-driven" target)
cor_un <- corMatrix(gls_un_sim$modelStruct$corStruct)[[1]]

# AR(1)-implied correlation matrix: rho^|j-k|
rho_hat <- as.numeric(coef(gls_ar1$modelStruct$corStruct, unconstrained = FALSE))
lags    <- abs(outer(seq_len(m), seq_len(m), "-"))
cor_ar1 <- rho_hat ^ lags

round(cor_un, 2)    # unstructured estimate

# REML-LRT: AR(1) (reduced) nested in UN (full).
# Both share the same maximal mean model, so REML logLik is comparable.
G2 <- 2 * (as.numeric(logLik(gls_un_sim)) - as.numeric(logLik(gls_ar1)))
df <- attr(logLik(gls_un_sim), "df") - attr(logLik(gls_ar1), "df")
p  <- pchisq(G2, df = df, lower.tail = FALSE)

data.frame(G2 = round(G2, 2), df = df, p_value = signif(p, 3))

# Fit three covariance choices with the same mean
gls_ind <- gls(y ~ group * time_f, data = long, method = "REML")

gls_ar1 <- gls(y ~ group * time_f, data = long,
               correlation = corAR1(form = ~ time_idx | id), method = "REML")

gls_un  <- gls(y ~ group * time_f, data = long,
               correlation = corSymm(value = start_vec_un, form = ~ time_idx | id),
               weights     = w_time,
               method = "REML")

# Helper: within-group change (week1 - baseline) via emmeans
get_change_tbl <- function(fit){
  emm_t_by_g <- emmeans(fit, ~ time_f | group)
  contr <- contrast(emm_t_by_g,
                    method = list("week1 - baseline" = c(-1, 1, 0, 0)))
  as.data.frame(summary(contr))[, c("group","estimate","SE","df","t.ratio","p.value")]
}

tab_ind <- get_change_tbl(gls_ind); tab_ind$model <- "Independence"
tab_ar1 <- get_change_tbl(gls_ar1); tab_ar1$model <- "AR(1)"
tab_un  <- get_change_tbl(gls_un ); tab_un$model  <- "UN (hetero)"

tab_all <- bind_rows(tab_ind, tab_ar1, tab_un) %>%
  select(model, group, estimate, SE, df, t.ratio, p.value)

tab_all

# UN with heterogeneous variances
gls_un <- gls(y ~ group * time_f, data = long,
              correlation = corSymm(value = start_vec_un, form = ~ time_idx | id),
              weights     = w_time,
              method      = "REML")

# Compound Symmetry (homoscedastic)
gls_cs <- gls(y ~ group * time_f, data = long,
              correlation = corCompSymm(form = ~ 1 | id),
              method      = "REML")

# AR(1) with equal spacing
gls_ar1 <- gls(y ~ group * time_f, data = long,
               correlation = corAR1(form = ~ time_idx | id),
               method      = "REML")

# Exponential (continuous-time)
gls_exp <- gls(y ~ group * time_f, data = long,
               correlation = corExp(form = ~ time_num | id),
               method      = "REML")

# Exponential with heterogeneous variances
gls_exp_het <- gls(y ~ group * time_f, data = long,
                   correlation = corExp(form = ~ time_num | id),
                   weights     = w_time,
                   method      = "REML")

# Exponential + nugget
gls_exp_nug <- gls(y ~ group * time_f, data = long,
                   correlation = corExp(form = ~ time_num | id, nugget = TRUE),
                   method      = "REML")

n <- 2:10
params_UN <- n*(n+1)/2
params_AR1 <- rep(2, length(n))
params_CS  <- rep(2, length(n))
df <- data.frame(n, UN=params_UN, AR1=params_AR1, CS=params_CS) |>
  tidyr::pivot_longer(-n, names_to="model", values_to="p")

ggplot(df, aes(n, p, color=model)) +
  geom_line() + geom_point() +
  labs(title="Number of covariance parameters vs visits",
       x="Number of visits (n)", y="Number of parameters") +
  theme_minimal()

N <- dplyr::n_distinct(long$id)
df2 <- df |>
  tidyr::pivot_wider(names_from = model, values_from = p) |>
  dplyr::mutate(N = N,
         N_per_param_UN = N/UN,
         N_per_param_AR1 = N/AR1,
         N_per_param_CS  = N/CS)
df2

library(reshape2)

# Correlation matrix from UN (first subject)
C_un_list <- corMatrix(gls_un$modelStruct$corStruct)
C_un <- if (is.list(C_un_list)) C_un_list[[1]] else C_un_list

dfC <- melt(C_un)
ggplot(dfC, aes(Var1, Var2, fill = value)) +
  geom_tile() + coord_equal() +
  scale_fill_gradient2(limits = c(-1,1)) +
  labs(title = "Implied correlation (UN)", x = "Visit index (1..4)", y = "Visit index (1..4)", fill = "rho") +
  theme_minimal()

# varIdent returns relative SD multipliers by factor level (base level=1)
w_coef <- coef(gls_un$modelStruct$varStruct, unconstrained = FALSE)
levs <- levels(long$time_f)
w_all <- setNames(rep(1, length(levs)), levs)
w_all[names(w_coef)] <- w_coef

sigma_eps <- sigma(gls_un)
sd_time <- sigma_eps * w_all

data.frame(time = names(sd_time), sd = as.numeric(sd_time)) %>%
  ggplot(aes(x = time, y = sd, group = 1)) +
  geom_line() + geom_point(size = 2) +
  labs(title = "Estimated residual SD by time (UN + varIdent)",
       x = "Time (weeks)", y = "SD") +
  theme_minimal()

summary(gls_cs)$tTable[1:6, , drop = FALSE]

summary(gls_ar1)$tTable[1:6, , drop = FALSE]

summary(gls_exp)$tTable[1:6, , drop = FALSE]

start_rho <- 0.4

gls_ar1_het <- gls(y ~ group * time_f, data = long,
      correlation = corAR1(value = start_rho, form = ~ time_idx | id),
      weights     = w_time,
      method      = "REML")

coef(gls_ar1_het$modelStruct$corStruct)

library(reshape2); library(dplyr); library(ggplot2)

get_cor_mat <- function(fit){
  M <- nlme::corMatrix(fit$modelStruct$corStruct)
  if (is.list(M)) M[[1]] else M
}

mats <- list(CS=gls_cs, AR1=gls_ar1, EXP=gls_exp, UN=gls_un) |>
  purrr::imap(~melt(get_cor_mat(.x)) |> mutate(model=.y))

bind_rows(mats) |>
  ggplot(aes(Var1, Var2, fill=value)) +
  geom_tile() + coord_equal() +
  scale_fill_gradient2(limits=c(-1,1)) +
  facet_wrap(~model, nrow=1) +
  labs(title="Implied correlation by Sigma", x="Visit index", y="Visit index", fill="rho") +
  theme_minimal()

mods <- list(
  UN_het   = gls_un,
  CS       = gls_cs,
  AR1      = gls_ar1,
  AR1het   = gls_ar1_het,
  EXP      = gls_exp,
  EXP_het  = gls_exp_het,
  EXP_nug  = gls_exp_nug
)

# Rank by AIC (FLW's recommended non-nested criterion); BIC for reference only.
data.frame(
  model  = names(mods),
  AIC    = sapply(mods, AIC),
  BIC    = sapply(mods, BIC),   # reference only; FLW p. 179 advises against for Sigma
  logLik = sapply(mods, logLik)
) |>
  dplyr::arrange(AIC)

# Independence (for boundary LRT demo)
gls_ind <- gls(y ~ group * time_f, data = long, method = "REML")

# CS, AR(1), Exponential, UN (hetero)
gls_cs  <- gls(y ~ group * time_f, data = long,
               correlation = corCompSymm(value = 0.2, form = ~ 1 | id), method = "REML")

gls_ar1 <- gls(y ~ group * time_f, data = long,
               correlation = corAR1(value = 0.2, form = ~ time_idx | id), method = "REML")

gls_exp <- gls(y ~ group * time_f, data = long,
               correlation = corExp(value = 0.2, form = ~ time_num | id), method = "REML")

gls_un  <- gls(y ~ group * time_f, data = long,
               correlation = corSymm(value = start_vec_un, form = ~ time_idx | id),
               weights     = w_time, method = "REML")

start_rho <- 0.4
gls_ar1_het <- gls(y ~ group * time_f, data = long,
                   correlation = corAR1(value = start_rho, form = ~ time_idx | id),
                   weights     = w_time, method = "REML")

mods <- list(
  IND     = gls_ind,
  CS      = gls_cs,
  AR1     = gls_ar1,
  AR1het  = gls_ar1_het,
  EXP     = gls_exp,
  UN_het  = gls_un
)

# AIC ranks non-nested candidates; for nested pairs use the REML-LRT.
# BIC shown for reference only (FLW p. 179 advises against it for Sigma).
data.frame(
  model  = names(mods),
  AIC    = sapply(mods, AIC),
  BIC    = sapply(mods, BIC),
  logLik = sapply(mods, logLik)
) |>
  dplyr::arrange(AIC)

# REML LRT: IND (reduced) vs CS read as a random intercept (full).
# Testing the random-intercept variance H0: sigma_b^2 = 0 (on the boundary).
LR <- 2 * (as.numeric(logLik(gls_cs)) - as.numeric(logLik(gls_ind)))
df <- 1  # one extra variance component
p_standard <- pchisq(LR, df = df, lower.tail = FALSE)           # naive
p_mixture  <- 0.5 * pchisq(LR, df = df, lower.tail = FALSE)     # correct boundary mixture

data.frame(LR = LR, df = df,
           p_naive_chisq = p_standard,
           p_mixture_50_50 = p_mixture)

# Fit random intercept model using lme()
lme_ri <- lme(y ~ group * time_f, random = ~ 1 | id, data = long, method = "REML")

# Compare fixed effects tables
summary(gls_cs)$tTable