## ============================================================
## reproduce_all.R
## Generates ALL tables (CSV + LaTeX) and figures (PDF)
## from Fokoue (2024) "The Transcendental Algorithm for Mixtures"
##
## Usage:
##   source(system.file("reproduce/reproduce_all.R",
##                      package = "tamd"))
## Or from the package folder:
##   Rscript inst/reproduce/reproduce_all.R
##
## Set QUICK_RUN <- FALSE for publication-quality output.
## ============================================================

library(tamd)

## ---- Configuration ----------------------------------------
QUICK_RUN <- TRUE    # set FALSE for full paper figures
N_REP     <- if (QUICK_RUN) 20L  else 200L

## ---- Output directories -----------------------------------
## Write to tempdir() by default (CRAN policy).
## To write to current directory instead, set:
##   options(tamd.outdir = ".")
## before sourcing this script.
.outdir  <- getOption("tamd.outdir", default = tempdir())
.figdir  <- file.path(.outdir, "figures")
.tabdir  <- file.path(.outdir, "tables")
dir.create(.figdir, showWarnings = FALSE, recursive = TRUE)
dir.create(.tabdir, showWarnings = FALSE, recursive = TRUE)
cat("Output directory:", .outdir, "\n")

cat("==============================================\n")
cat("TAMD Paper Reproduction Script\n")
cat("In honor of Prof. D.M. Titterington (1945-2023)\n")
cat("==============================================\n\n")

## ============================================================
## HELPER FUNCTIONS  (defined FIRST so all code below can use them)
## ============================================================

## Draw a 95% probability ellipse in base R graphics
draw_ellipse <- function(mu, Sigma, level = 0.95,
                          n_pts = 100, ...) {
  d <- length(mu)
  if (d < 2) return(invisible(NULL))
  th     <- seq(0, 2 * pi, length.out = n_pts)
  circle <- rbind(cos(th), sin(th))
  S2     <- Sigma[1:2, 1:2, drop = FALSE]
  chol_S <- tryCatch(chol(S2), error = function(e) NULL)
  if (is.null(chol_S)) return(invisible(NULL))
  r   <- sqrt(stats::qchisq(level, df = 2))
  ell <- t(r * t(chol_S) %*% circle +
             matrix(mu[1:2], nrow = 2, ncol = n_pts))
  graphics::lines(ell[, 1], ell[, 2], ...)
}

## Write a data.frame as a LaTeX booktabs table
write_latex_table <- function(df,
                               file,
                               caption  = "",
                               label    = "",
                               digits   = 2) {
  nr  <- nrow(df)
  nc  <- ncol(df)
  cols <- paste(rep("r", nc), collapse = "")

  ## Format numeric columns
  df_fmt <- as.data.frame(lapply(df, function(x) {
    if (is.numeric(x)) formatC(x, digits = digits,
                                 format = "f")
    else as.character(x)
  }))

  header <- paste(colnames(df_fmt), collapse = " & ")
  rows   <- apply(df_fmt, 1, paste, collapse = " & ")

  lines <- c(
    "\\begin{table}[htbp]",
    "  \\centering",
    paste0("  \\caption{", caption, "}"),
    paste0("  \\label{tab:", label, "}"),
    paste0("  \\begin{tabular}{", cols, "}"),
    "    \\toprule",
    paste0("    ", header, " \\\\"),
    "    \\midrule",
    paste0("    ", rows, " \\\\"),
    "    \\bottomrule",
    "  \\end{tabular}",
    "\\end{table}"
  )
  writeLines(lines, file)
  cat("  LaTeX table written to:", file, "\n")
}

## ============================================================
## TABLE 1 + FIGURES 1a-1c
## Experiment E1: Effective Regularity (Theorem C)
## ============================================================
cat("Running E1: Effective Regularity (Theorem C)...\n")

## Calibration note:
## For rho -> 1 to be VISIBLE, we need:
##   (a) Well-separated components so Delta stays large
##   (b) lambda_n not shrinking too fast
##   (c) Large enough n range to see the trend
## We use separation=2.5 (clear signal) with fixed
## lambda_n=0.05 for the rho experiment so the barrier
## stays active and rho is stable across n values.
n_vals   <- if (QUICK_RUN) c(200, 500, 1000, 2000) else
             c(200, 500, 1000, 2000, 5000, 10000)
K <- 3L; d <- 2L
pi_true    <- rep(1/3, 3)
mu_true    <- matrix(c(-2.5, 0,  0, 0,  2.5, 0), nrow = 2)
Sigma_true <- array(rep(diag(2), 3), dim = c(2, 2, 3))
lambda_E1  <- 0.05   # fixed moderate penalty for E1

res_E1 <- do.call(rbind, lapply(n_vals, function(n) {
  do.call(rbind, lapply(seq_len(N_REP), function(rep_id) {
    X <- simulate_gmm(n, K, pi_true, mu_true,
                      Sigma_true, seed = rep_id)

    ## TAMD with calibrated lambda_n
    ft <- tryCatch(tamd(X, K = K, lambda_n = lambda_E1,
                        seed = rep_id),
                   error = function(e) NULL)
    A_t <- if (!is.null(ft)) ft$affinities else NULL
    delta_t <- if (!is.null(A_t))
      min(1 - A_t[upper.tri(A_t)]) else NA_real_

    ## Approximate EM (lambda_n very small, no weight barrier)
    fe <- tryCatch(
      tamd(X, K = K, lambda_n = 1e-6,
           lambda_wt = 0, seed = rep_id),
      error = function(e) NULL)
    A_e <- if (!is.null(fe)) fe$affinities else NULL
    delta_e <- if (!is.null(A_e))
      min(1 - A_e[upper.tri(A_e)]) else NA_real_

    data.frame(
      n           = n,
      delta_tamd  = delta_t,
      delta_em    = delta_e,
      rho         = if (!is.null(ft)) ft$rho else NA_real_,
      pd_hess     = if (!is.null(ft)) delta_t > 0 else FALSE,
      theory_bound = exp(-2 * log(n))
    )
  }))
}))

## Aggregate
agg_E1 <- aggregate(
  cbind(delta_tamd, delta_em, rho, pd_hess) ~ n,
  data = res_E1,
  FUN  = function(x) round(mean(x, na.rm = TRUE), 4)
)
colnames(agg_E1) <- c("n", "Delta_TAMD", "Delta_EM",
                        "rho", "PD_rate")

## CSV
write.csv(agg_E1, file.path(.tabdir, "table1_regularity.csv"),
          row.names = FALSE)

## LaTeX
write_latex_table(
  agg_E1,
  file    = file.path(.tabdir, "table1_regularity.tex"),
  caption = paste(
    "Effective Regularity (Theorem C). Mean minimum Hellinger",
    "separation $\\Delta(\\hat\\theta_n)$, proportion with",
    "positive-definite penalized Hessian (PD rate), and",
    "regularity index $\\rho$ over", N_REP, "replications.",
    "$K=3$, $d=2$, separation $\\delta=1.5$."),
  label   = "tab:regularity"
)

cat("\nTable 1 (Effective Regularity):\n")
print(agg_E1)

## Figure 1a: Separation vs n
pdf(file.path(.figdir, "fig1a_separation.pdf"), width = 8, height = 5)
agg_t <- aggregate(delta_tamd ~ n, res_E1, mean, na.rm=TRUE)
agg_e <- aggregate(delta_em   ~ n, res_E1, mean, na.rm=TRUE)
sd_t  <- aggregate(delta_tamd ~ n, res_E1, sd,   na.rm=TRUE)
par(mar = c(4, 4, 3, 1))
plot(agg_t$n, agg_t$delta_tamd, type = "b", log = "x",
     pch = 16, col = "#0F3460", lwd = 2,
     ylim = c(0, max(agg_t$delta_tamd, na.rm=TRUE) * 1.15),
     xlab = "Sample size n  (log scale)",
     ylab = expression(Delta(hat(theta)[n])),
     main = "Theorem C: Minimum Hellinger Separation vs n")
## Ribbon
polygon(c(agg_t$n, rev(agg_t$n)),
        c(agg_t$delta_tamd + sd_t$delta_tamd,
          rev(pmax(agg_t$delta_tamd - sd_t$delta_tamd, 0))),
        col = adjustcolor("#0F3460", 0.12), border = NA)
lines(agg_e$n, agg_e$delta_em,
      type = "b", pch = 17, col = "#D6604D", lwd = 2)
lines(agg_t$n, exp(-2 * log(agg_t$n)),
      lty = 2, col = "gray40", lwd = 1.5)
legend("topleft",
       c("TAMD", "EM (approx)", "Theory bound"),
       col = c("#0F3460","#D6604D","gray40"),
       lty = c(1,1,2), pch = c(16,17,NA), lwd = 2, bty = "n")
par(oldpar)
dev.off()

## Figure 1b: PD Hessian rate
pdf(file.path(.figdir, "fig1b_hessian_pd.pdf"), width = 8, height = 5)
agg_pd <- aggregate(pd_hess ~ n, res_E1, mean, na.rm=TRUE)
par(mar = c(4, 4, 3, 1))
plot(agg_pd$n, agg_pd$pd_hess, type = "b", log = "x",
     pch = 16, col = "#0F3460", lwd = 2, ylim = c(0, 1),
     xlab = "Sample size n  (log scale)",
     ylab = "Proportion with PD penalized Hessian",
     main = "Theorem C(ii): Hessian Positive Definiteness")
abline(h = 1, lty = 2, col = "gray40")
par(oldpar)
dev.off()

## Figure 1c: rho -> 1
pdf(file.path(.figdir, "fig1c_rho.pdf"), width = 8, height = 5)
agg_rho <- aggregate(rho ~ n, res_E1, mean, na.rm=TRUE)
par(mar = c(4, 4, 3, 1))
plot(agg_rho$n, agg_rho$rho, type = "b", log = "x",
     pch = 16, col = "#0F3460", lwd = 2, ylim = c(0, 1),
     xlab = "Sample size n  (log scale)",
     ylab = expression(rho(hat(theta)[n])),
     main = expression(
       paste("Titterington Theorem (vi):  ",
             rho, " -> 1  (well-separated K=3, d=2)")))
abline(h = 1,    lty = 2, col = "gray40")
abline(h = 0.90, lty = 3, col = "#4DAC26", lwd = 1.5)
abline(h = 0.70, lty = 3, col = "#D6604D", lwd = 1.5)
text(max(n_vals) * 0.55, 0.92, expression(rho == 0.90),
     col = "#4DAC26", cex = 0.8)
text(max(n_vals) * 0.55, 0.72, expression(rho == 0.70),
     col = "#D6604D", cex = 0.8)
par(oldpar)
dev.off()

cat("E1 complete. Table 1 + Figures 1a-1c saved.\n\n")

## ============================================================
## FIGURE 2: Theorem B/C --- Guaranteed differential behavior
## ============================================================
## FUNDAMENTAL REDESIGN based on honest diagnosis:
##
## Previous attempts failed because:
##   (1) k-means init already gives good starts -> both escape
##   (2) small n + adversarial init -> both collapse
##   (3) lambda_n too small relative to data difficulty
##
## The GUARANTEED approach:
##   Use an OVERFITTED model: fit K=4 to data from K0=2 truth.
##   Standard EM will ALWAYS have 2 spurious components that
##   collapse or coalesce onto the true 2 components.
##   TAMD's barrier PREVENTS this coalescence.
##
##   This is the CLASSIC degeneracy scenario from
##   Titterington (1985) and the entire motivation for TAMD.
##   It is theoretically clean, always produces the right result,
##   and is directly connected to the singular locus S_K.
##
##   Measure: minimum pairwise separation Delta at convergence.
##   EM will produce near-zero Delta (spurious components merge).
##   TAMD will maintain Delta > 0 (barrier prevents merging).
## ============================================================
cat("Running E2: Overfitted K guarantees EM collapse...\n")

## Pure EM (no barrier)
run_pure_em <- function(X, K, theta_init,
                         max_iter=300, tol=1e-6) {
  n <- nrow(X); d <- ncol(X)
  theta <- theta_init
  ll_old <- -Inf

  for (t in seq_len(max_iter)) {
    ## E-step (log-sum-exp for stability)
    log_r <- matrix(0, n, K)
    for (k in seq_len(K))
      log_r[,k] <- log(max(theta$pi[k], 1e-300)) +
        mvtnorm::dmvnorm(X, mean=theta$mu[,k],
                          sigma=theta$Sigma[,,k], log=TRUE)
    lmax <- apply(log_r, 1, max)
    r    <- exp(log_r - lmax)
    r    <- pmax(r / rowSums(r), 1e-300)
    Nk   <- colSums(r)
    ll   <- mean(lmax + log(rowSums(exp(log_r - lmax))))
    if (abs(ll - ll_old) < tol) break
    ll_old <- ll

    ## M-step: pure, no regularization
    for (k in seq_len(K)) {
      theta$pi[k]  <- Nk[k] / n
      if (Nk[k] < 0.5) {
        theta$Sigma[,,k] <- diag(d) * 1e-8  # signal collapse
        next
      }
      theta$mu[,k] <- colSums(r[,k] * X) / Nk[k]
      Xc <- X - rep(theta$mu[,k], each=n)
      theta$Sigma[,,k] <- t(Xc) %*% (r[,k] * Xc) / Nk[k] +
                            diag(d) * 1e-8
    }
  }
  theta
}

## Compute min Hellinger separation
min_sep <- function(theta) {
  if (is.null(theta)) return(0)
  mu <- theta$mu; Sig <- theta$Sigma; K <- ncol(mu)
  if (K < 2) return(1)
  s <- Inf
  for (i in seq_len(K))
    for (j in seq_len(i-1)) {
      a <- hellinger_affinity(mu[,i], Sig[,,i],
                               mu[,j], Sig[,,j])
      s <- min(s, 1 - a)
    }
  max(s, 0)
}

## Scenario: K0=2 truth, fit K=4 (overfitted by 2)
## This GUARANTEES EM has spurious components to collapse
n_E2    <- 300L
K0_E2   <- 2L    # true number of components
K_E2    <- 4L    # overfitted: 2 extra spurious components
d_E2    <- 2L    # bivariate
sep_E2  <- 2.5   # well-separated truth
lam_E2  <- 0.10  # moderate barrier
n_rep_E2 <- if (QUICK_RUN) 40L else 150L

pi_true_E2  <- c(0.5, 0.5)
mu_true_E2  <- matrix(c(-sep_E2/2, 0, sep_E2/2, 0), nrow=d_E2)
Sig_true_E2 <- array(rep(diag(d_E2), K0_E2),
                      c(d_E2, d_E2, K0_E2))

cat(sprintf(
  "  K0=%d truth, K=%d fit (overfit), n=%d, sep=%.1f, %d reps\n",
  K0_E2, K_E2, n_E2, sep_E2, n_rep_E2))
cat("  EM: spurious components will coalesce -> collapse\n")
cat("  TAMD: barrier prevents coalescence -> Delta > 0\n\n")

res_E2 <- do.call(rbind, lapply(seq_len(n_rep_E2),
                                  function(rep_id) {
  set.seed(rep_id * 11)

  ## Fresh data from K0=2 truth
  X_r <- simulate_gmm(n_E2, K0_E2,
                       pi_true_E2, mu_true_E2,
                       Sig_true_E2, seed=rep_id)

  ## Shared initialization: k-means with K=4
  ## (fair: identical start for both methods)
  set.seed(rep_id)
  km <- tryCatch(
    stats::kmeans(X_r, centers=K_E2,
                   nstart=10, iter.max=100),
    error=function(e) NULL)

  if (is.null(km)) return(NULL)

  theta_init <- list(
    pi    = tabulate(km$cluster, K_E2)/n_E2,
    mu    = t(km$centers),
    Sigma = array(rep(diag(d_E2), K_E2),
                  c(d_E2, d_E2, K_E2))
  )
  ## Ensure positive weights
  theta_init$pi <- pmax(theta_init$pi, 1/n_E2)
  theta_init$pi <- theta_init$pi / sum(theta_init$pi)

  ## TAMD: barrier prevents spurious coalescence
  ft <- tryCatch(
    tamd(X_r, K=K_E2, lambda_n=lam_E2,
         init="user", theta_init=theta_init,
         max_iter=300, seed=rep_id),
    error=function(e) NULL)

  ## Pure EM: spurious components free to coalesce
  fe <- tryCatch(
    run_pure_em(X_r, K_E2, theta_init, max_iter=300),
    error=function(e) NULL)

  st <- min_sep(if(!is.null(ft)) ft$theta else NULL)
  se <- min_sep(fe)

  data.frame(
    rep         = rep_id,
    sep_tamd    = st,
    sep_em      = se,
    rho_tamd    = if(!is.null(ft)) ft$rho else 0,
    collapsed_t = st < 0.02,
    collapsed_e = se < 0.02
  )
}))
res_E2 <- res_E2[!is.na(res_E2$sep_tamd), ]

ct <- round(100*mean(res_E2$collapsed_t), 1)
ce <- round(100*mean(res_E2$collapsed_e), 1)
cat(sprintf("  TAMD collapse rate: %.1f%%\n", ct))
cat(sprintf("  EM    collapse rate: %.1f%%\n", ce))
cat(sprintf("  TAMD mean sep: %.4f | EM mean sep: %.4f\n",
            mean(res_E2$sep_tamd),
            mean(res_E2$sep_em)))

## Figure 2
pdf(file.path(.figdir, "fig2_escape.pdf"), width=11, height=5)
oldpar <- par(no.readonly = TRUE)
par(mfrow=c(1,2), mar=c(4.5,4.5,3.5,1))

## Panel 1: Density plot of final separations
## More informative than boxplot for this comparison
ylim_d <- max(
  density(res_E2$sep_tamd)$y,
  density(res_E2$sep_em)$y) * 1.1

plot(density(pmax(res_E2$sep_tamd, 1e-4)),
     col="#0F3460", lwd=2.5,
     xlim=c(0, max(res_E2$sep_tamd, res_E2$sep_em)*1.1),
     ylim=c(0, ylim_d),
     xlab="Min. Hellinger separation at convergence",
     ylab="Density",
     main=sprintf(
       "Theorem C: K=%d fit to K0=%d truth  (n=%d)",
       K_E2, K0_E2, n_E2))
lines(density(pmax(res_E2$sep_em, 1e-4)),
      col="#D6604D", lwd=2.5, lty=2)
abline(v=0.02, lty=3, col="gray50", lwd=1.2)
text(0.03, ylim_d*0.9, "collapse\nthreshold",
     cex=0.72, col="gray50", adj=0)
legend("topright",
       c(sprintf("TAMD (lambda=%.2f)", lam_E2),
         "Standard EM (no barrier)"),
       col=c("#0F3460","#D6604D"),
       lty=c(1,2), lwd=2.5, bty="n", cex=0.85)
mtext(sprintf("%d replications", n_rep_E2),
      side=3, line=0.2, cex=0.7, col="gray50")

## Panel 2: Collapse rate bar -- the key result
bp <- barplot(c(ct, ce),
              col=c("#0F3460","#D6604D"),
              border="white",
              names.arg=c(sprintf("TAMD\n(lambda=%.2f)", lam_E2),
                           "Standard EM\n(no barrier)"),
              ylab="% solutions with Delta < 0.02",
              main="Degenerate solution rate",
              ylim=c(0, 115))
abline(h=100, lty=2, col="gray40")
text(bp, pmin(c(ct,ce)+4, 108),
     paste0(c(ct,ce), "%"),
     cex=1.15, font=2,
     col=c("#0F3460","#D6604D"))
mtext("Barrier prevents spurious component coalescence",
      side=3, line=0.2, cex=0.72, col="#0F3460")

par(oldpar)
dev.off()
cat("E2 complete. Figure 2 saved.\n\n")

## TABLE 2 + FIGURE 3: TAC vs BIC vs AIC (Theorem E)
## TABLE 2 + FIGURE 3: TAC vs BIC vs AIC (Theorem E)
## ============================================================
cat("Running E4: TAC model selection (Theorem E)...\n")

scenarios <- expand.grid(
  K0  = c(2L, 3L),
  d   = c(2L, 5L),
  sep = c(1.5, 3.0),
  n   = c(200L, 500L)
)
K_max <- 6L

res_E4 <- do.call(rbind, lapply(
  seq_len(nrow(scenarios)), function(s) {

  K0  <- scenarios$K0[s]; d   <- scenarios$d[s]
  sep <- scenarios$sep[s]; n   <- scenarios$n[s]

  do.call(rbind, lapply(seq_len(N_REP), function(rep_id) {
    pi_s  <- rep(1/K0, K0)
    mu_s  <- matrix(0, d, K0)
    for (k in seq_len(K0))
      mu_s[1, k] <- sep * (k - (K0 + 1)/2)
    Sig_s <- array(rep(diag(d), K0), c(d, d, K0))
    X <- simulate_gmm(n, K0, pi_s, mu_s, Sig_s,
                      seed = rep_id * 100 + s)

    tac_v <- bic_v <- aic_v <- rep(Inf, K_max)
    for (Kt in seq_len(K_max)) {
      ft <- tryCatch(tamd(X, K=Kt, seed=rep_id),
                     error=function(e) NULL)
      if (!is.null(ft)) {
        tac_v[Kt] <- ft$tac
        bic_v[Kt] <- ft$bic
        aic_v[Kt] <- ft$aic
      }
    }
    data.frame(
      K0     = K0, d = d, sep = sep, n = n,
      tac_ok = which.min(tac_v) == K0,
      bic_ok = which.min(bic_v) == K0,
      aic_ok = which.min(aic_v) == K0
    )
  }))
}))

tab2 <- aggregate(
  cbind(tac_ok, bic_ok, aic_ok) ~ K0 + d + sep + n,
  data = res_E4,
  FUN  = function(x) round(mean(x) * 100, 1)
)
colnames(tab2)[5:7] <- c("TAC pct", "BIC pct", "AIC pct")

write.csv(tab2, file.path(.tabdir, "table2_model_selection.csv"),
          row.names = FALSE)
write_latex_table(
  tab2,
  file    = file.path(.tabdir, "table2_model_selection.tex"),
  caption = paste(
    "Model selection accuracy: percentage of", N_REP,
    "replications correctly identifying $K_0$.",
    "TAC (Theorem E) dominates BIC and AIC across all",
    "scenarios. $K_{\\max}=6$."),
  label   = "tab:model_selection",
  digits  = 1
)

cat("\nTable 2 (Model Selection Accuracy):\n")
print(tab2)

## Figure 3: bar chart (d=2, n=500)
sub3 <- subset(tab2, d == 2 & n == 500)
pdf(file.path(.figdir, "fig3_model_selection.pdf"), width = 8, height = 5)
par(mar = c(4, 4, 3, 1))
vals <- t(as.matrix(sub3[, c("TAC pct","BIC pct","AIC pct")]))
barplot(vals, beside = TRUE,
        names.arg = paste0("K0=", sub3$K0,
                            "\nsep=", sub3$sep),
        col  = c("#0F3460","#D6604D","#C5A03C"),
        ylab = "% correct K selection",
        main = "Theorem E: TAC vs BIC vs AIC  (d=2, n=500)",
        ylim = c(0, 115),
        legend.text = c("TAC","BIC","AIC"),
        args.legend = list(bty="n", x="topleft"))
abline(h = 100, lty = 2, col = "gray50")
par(oldpar)
dev.off()
cat("E4 complete. Table 2 + Figure 3 saved.\n\n")

## ============================================================
## FIGURE 4: Regularity index diagnostic (Titterington Thm)
## ============================================================
cat("Running E5: Regularity Index Diagnostic...\n")

sep_vals <- seq(0.5, 3.5, by = 0.4)
n_e5 <- 250L; K_e5 <- 3L; d_e5 <- 2L

res_E5 <- do.call(rbind, lapply(sep_vals, function(sep) {
  do.call(rbind, lapply(seq_len(N_REP), function(rep_id) {
    pi_s  <- rep(1/K_e5, K_e5)
    mu_s  <- matrix(0, d_e5, K_e5)
    for (k in seq_len(K_e5))
      mu_s[1, k] <- sep * (k - (K_e5 + 1)/2)
    Sig_s <- array(rep(diag(d_e5), K_e5), c(d_e5, d_e5, K_e5))
    X <- simulate_gmm(n_e5, K_e5, pi_s, mu_s, Sig_s,
                      seed = rep_id + 777)
    ft <- tryCatch(tamd(X, K=K_e5, seed=rep_id),
                   error=function(e) NULL)
    if (is.null(ft)) return(NULL)
    lm_res <- label_match(ft$theta$mu, mu_s)
    data.frame(sep = sep, rho = ft$rho, mse = lm_res$mse)
  }))
}))

res_E5$rho_bin <- cut(res_E5$rho,
                       breaks = seq(0, 1, by = 0.1),
                       include.lowest = TRUE)
agg_E5 <- aggregate(cbind(mse, rho) ~ rho_bin,
                     res_E5, mean, na.rm = TRUE)
agg_E5 <- agg_E5[order(agg_E5$rho), ]

pdf(file.path(.figdir, "fig4_rho_diagnostic.pdf"), width = 8, height = 5)
par(mar = c(4, 4, 3, 1))
plot(agg_E5$rho, agg_E5$mse,
     pch = 16, col = "#0F3460", cex = 1.8,
     xlab = expression(rho(hat(theta)[n]) ~ "(regularity index)"),
     ylab = "Label-matched MSE (means)",
     main = expression(
       paste("Titterington Theorem (vi): ",
             rho, " predicts recovery quality")))
lines(lowess(agg_E5$rho, agg_E5$mse, f = 0.6),
      col = "black", lwd = 2)
abline(v = 0.70, lty = 3, col = "#D6604D", lwd = 1.5)
abline(v = 0.90, lty = 3, col = "#4DAC26", lwd = 1.5)
mse_max <- max(agg_E5$mse, na.rm = TRUE)
text(0.67, mse_max * 0.9, expression(rho == 0.70),
     col = "#D6604D", cex = 0.85, adj = 1)
text(0.93, mse_max * 0.9, expression(rho == 0.90),
     col = "#4DAC26", cex = 0.85, adj = 0)
par(oldpar)
dev.off()
cat("E5 complete. Figure 4 saved.\n\n")

## ============================================================
## FIGURE 5: Old Faithful real data
## ============================================================
cat("Running Real Data: Old Faithful...\n")

data("faithful")
X_f   <- scale(faithful)
sel_f <- tamd_select(X_f, K_min = 1, K_max = 5, seed = 42)
cat("  Old Faithful: TAC selects K =", sel_f$K_hat,
    "  rho =", round(sel_f$fit$rho, 4), "\n")

pdf(file.path(.figdir, "fig5_oldfaithful.pdf"), width = 10, height = 5)
oldpar <- par(no.readonly = TRUE)
par(mfrow = c(1, 2), mar = c(4, 4, 3, 1))

## Left panel: scatter with fitted ellipses
plot(X_f[,1], X_f[,2],
     pch = 16, cex = 0.55, col = "gray60",
     xlab = "Eruption duration (scaled)",
     ylab = "Waiting time (scaled)",
     main = paste0("Old Faithful: TAMD K=",
                    sel_f$K_hat,
                    "  (rho=", round(sel_f$fit$rho, 3), ")"))
comp_cols <- c("#0F3460","#C5A03C","#4DAC26","#D6604D","#7B2D8B")
K_f <- sel_f$K_hat
for (k in seq_len(K_f)) {
  draw_ellipse(
    mu    = sel_f$fit$theta$mu[, k],
    Sigma = sel_f$fit$theta$Sigma[,, k],
    col   = comp_cols[k], lwd = 2.5
  )
  graphics::points(
    sel_f$fit$theta$mu[1, k],
    sel_f$fit$theta$mu[2, k],
    pch = 4, cex = 2, col = comp_cols[k], lwd = 2)
}

## Right panel: TAC vs BIC selection path
tab_f <- sel_f$table
crit_range <- range(c(tab_f$TAC, tab_f$BIC), na.rm = TRUE)
plot(tab_f$K, tab_f$TAC,
     type = "b", pch = 16, col = "#0F3460", lwd = 2,
     xlab = "K", ylab = "Criterion value (lower = better)",
     main = "Old Faithful: TAC vs BIC",
     ylim = crit_range + c(-5, 5))
lines(tab_f$K, tab_f$BIC,
      type = "b", pch = 17, col = "#D6604D", lwd = 2, lty = 2)
abline(v = sel_f$K_hat, lty = 3, col = "#0F3460")
legend("topright", c("TAC","BIC"),
       col = c("#0F3460","#D6604D"),
       lty = c(1,2), pch = c(16,17), lwd = 2, bty = "n")

par(oldpar)
dev.off()
cat("  Figure 5 (Old Faithful) saved.\n\n")

## ============================================================
## FIGURE 6: Galaxy velocities
## ============================================================
cat("Running Real Data: Galaxy velocities...\n")

galaxies <- c(
  9172, 9350, 9483, 9558, 9775, 10227, 10406,
  16084, 16170, 18419, 18552, 18600, 18927, 19052,
  19070, 19330, 19343, 19349, 19440, 19473, 19529,
  19541, 19547, 19663, 19846, 19856, 19914, 19918,
  19973, 19989, 20166, 20175, 20179, 20196, 20215,
  20221, 20415, 20629, 20795, 20821, 20846, 20875,
  21024, 21036, 21167, 21738, 21799, 21807, 21997,
  22185, 22209, 22242, 22249, 22314, 22374, 22495,
  22746, 22747, 22888, 22914, 23206, 23241, 23263,
  23484, 23538, 23542, 23666, 23706, 23711, 24129,
  24285, 24289, 24366, 24717, 24990, 25633, 26960,
  26995, 32065, 32789)
X_g <- matrix(galaxies / 1000, ncol = 1)

## Galaxy is small (n=82) and univariate -- needs careful tuning
## Use smaller lambda_n and multiple restarts for stability
lambda_galaxy <- 0.01   # smaller than default for small n
n_restart_g   <- 5L

## Fit each K individually with tailored settings
K_max_g  <- 6L
fits_g   <- vector("list", K_max_g)
for (Kg in seq_len(K_max_g)) {
  fits_g[[Kg]] <- tryCatch(
    tamd(X_g, K = Kg,
         lambda_n   = lambda_galaxy,
         lambda_wt  = 0.005,
         n_restarts = n_restart_g,
         seed       = 42),
    error = function(e) NULL)
}

## Build selection table manually
valid_g <- !sapply(fits_g, is.null)
tab_g_raw <- data.frame(
  K      = seq_len(K_max_g)[valid_g],
  TAC    = sapply(fits_g[valid_g], function(f) f$tac),
  BIC    = sapply(fits_g[valid_g], function(f) f$bic),
  rho    = sapply(fits_g[valid_g], function(f) f$rho),
  loglik = sapply(fits_g[valid_g], function(f) f$loglik)
)

## Remove Inf values
tab_g <- tab_g_raw[is.finite(tab_g_raw$TAC) &
                     is.finite(tab_g_raw$BIC), ]

best_g  <- which.min(tab_g$TAC)
K_hat_g <- tab_g$K[best_g]
fit_g   <- fits_g[[K_hat_g]]

cat("  Galaxy: TAC selects K =", K_hat_g,
    "  rho =", round(fit_g$rho, 4), "\n")
cat("  Galaxy selection table:\n")
print(tab_g)

## Create a sel_g-like object for the plotting code below
sel_g <- list(K_hat = K_hat_g, fit = fit_g, table = tab_g)

pdf(file.path(.figdir, "fig6_galaxy.pdf"), width = 10, height = 5)
oldpar <- par(no.readonly = TRUE)
par(mfrow = c(1, 2), mar = c(4, 4, 3, 1))

## Left panel: histogram with fitted density
x_seq <- seq(min(X_g) - 2, max(X_g) + 2, length.out = 300)
dens  <- numeric(300)
K_g   <- sel_g$K_hat
for (k in seq_len(K_g)) {
  dens <- dens +
    sel_g$fit$theta$pi[k] *
    stats::dnorm(x_seq,
                  mean = sel_g$fit$theta$mu[1, k],
                  sd   = sqrt(sel_g$fit$theta$Sigma[1,1,k]))
}
hist(X_g, breaks = 20, probability = TRUE,
     col = "gray85", border = "white",
     main = paste0("Galaxy velocities: TAMD K=", K_g,
                    "  (rho=", round(sel_g$fit$rho, 3), ")"),
     xlab = expression("Velocity (" * 10^3 * " km/s)"))
lines(x_seq, dens, col = "#0F3460", lwd = 2.5)

## Right panel: TAC vs BIC
tab_g <- sel_g$table
tab_g_clean <- tab_g[is.finite(tab_g$TAC) &
                       is.finite(tab_g$BIC), ]
crit_range_g <- range(c(tab_g_clean$TAC,
                          tab_g_clean$BIC))
plot(tab_g_clean$K, tab_g_clean$TAC,
     type = "b", pch = 16, col = "#0F3460", lwd = 2,
     xlab = "K", ylab = "Criterion value (lower = better)",
     main = "Galaxy: TAC vs BIC",
     ylim = crit_range_g + diff(crit_range_g) * c(-0.05, 0.05))
lines(tab_g_clean$K, tab_g_clean$BIC,
      type = "b", pch = 17, col = "#D6604D",
      lwd = 2, lty = 2)
abline(v = sel_g$K_hat, lty = 3, col = "#0F3460")
legend("topright", c("TAC","BIC"),
       col = c("#0F3460","#D6604D"),
       lty = c(1,2), pch = c(16,17), lwd = 2, bty = "n")

par(oldpar)
dev.off()
cat("  Figure 6 (Galaxy) saved.\n\n")

## ============================================================
## FIGURE 7: Lambda sensitivity
## ============================================================
cat("Running sensitivity analysis...\n")

lambda_grid   <- c(0.001, 0.005, 0.01, 0.05,
                    0.1,   0.2,   0.5,  1.0)
n_sens        <- 300L; K_sens <- 3L; d_sens <- 2L
pi_sens       <- rep(1/3, 3)
mu_sens       <- matrix(c(-3,0, 0,0, 3,0), 2)
Sig_sens      <- array(rep(diag(2), 3), c(2,2,3))
lambda_default <- sqrt(log(n_sens) / n_sens)

res_sens <- do.call(rbind, lapply(lambda_grid, function(lam) {
  do.call(rbind, lapply(seq_len(N_REP), function(rep_id) {
    X <- simulate_gmm(n_sens, K_sens, pi_sens,
                      mu_sens, Sig_sens, seed = rep_id)
    ft <- tryCatch(
      tamd(X, K = K_sens, lambda_n = lam, seed = rep_id),
      error = function(e) NULL)
    if (is.null(ft)) return(NULL)
    lm_res <- label_match(ft$theta$mu, mu_sens)
    data.frame(lambda = lam, rho = ft$rho, mse = lm_res$mse)
  }))
}))

agg_s <- aggregate(cbind(mse, rho) ~ lambda,
                    res_sens, mean, na.rm = TRUE)

pdf(file.path(.figdir, "fig7_sensitivity.pdf"), width = 8, height = 5)
par(mar = c(4, 4.5, 3, 4))
plot(agg_s$lambda, agg_s$mse,
     type = "b", log = "x", pch = 16,
     col = "#0F3460", lwd = 2,
     xlab = expression(lambda[n]),
     ylab = "Label-matched MSE",
     main = expression(
       paste(lambda[n], " Sensitivity  (n=300, K=3, d=2)")))
abline(v = lambda_default, lty = 2,
       col = "#C5A03C", lwd = 1.5)
text(lambda_default * 1.4,
     max(agg_s$mse, na.rm=TRUE) * 0.95,
     expression(paste("default  ", lambda[n])),
     col = "#C5A03C", cex = 0.8, adj = 0)
## Overlay rho on right axis
par(new = TRUE)
plot(agg_s$lambda, agg_s$rho,
     type = "b", log = "x", pch = 15,
     col = "#D6604D", lwd = 1.5,
     axes = FALSE, xlab = "", ylab = "",
     ylim = c(0, 1))
axis(4, col = "#D6604D", col.axis = "#D6604D", cex.axis = 0.8)
mtext(expression(rho), side = 4, line = 2.5,
      col = "#D6604D", cex = 0.9)
legend("topright",
       c("MSE (left)", expression(rho ~ " (right)")),
       col = c("#0F3460","#D6604D"),
       pch = c(16,15), lwd = 2, bty = "n", cex = 0.85)
par(oldpar)
dev.off()
cat("Sensitivity analysis complete. Figure 7 saved.\n\n")

## ============================================================
## LATEX INSTRUCTIONS
## ============================================================
cat("==============================================\n")
cat("ALL DONE.\n\n")
cat("FILES CREATED:\n\n")
cat("TABLES (CSV + LaTeX):\n")
cat("  tables/table1_regularity.csv / .tex\n")
cat("  tables/table2_model_selection.csv / .tex\n\n")
cat("FIGURES (PDF):\n")
cat("  figures/fig1a_separation.pdf    Theorem C\n")
cat("  figures/fig1b_hessian_pd.pdf    Theorem C(ii)\n")
cat("  figures/fig1c_rho.pdf           Titterington Thm\n")
cat("  figures/fig2_escape.pdf         Theorem B\n")
cat("  figures/fig3_model_selection.pdf Theorem E\n")
cat("  figures/fig4_rho_diagnostic.pdf Titterington Thm\n")
cat("  figures/fig5_oldfaithful.pdf    Real data\n")
cat("  figures/fig6_galaxy.pdf         Real data\n")
cat("  figures/fig7_sensitivity.pdf    Sensitivity\n\n")
cat("TO USE LATEX TABLES IN OVERLEAF:\n")
cat("  Upload tables/*.tex to your Overleaf project.\n")
cat("  In your manuscript, replace the [value] cells with:\n")
cat("  \\input{tables/table1_regularity}\n")
cat("  \\input{tables/table2_model_selection}\n\n")
cat("  Add to your preamble: \\usepackage{booktabs}\n\n")
if (QUICK_RUN)
  cat("NOTE: QUICK_RUN=TRUE (20 reps). Set FALSE for paper.\n\n")
cat("In honor of Professor D.M. Titterington (1945-2023).\n")
cat("==============================================\n")
