## ----setup, include = FALSE---------------------------------------------------
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  fig.width = 7,
  fig.height = 5,
  fig.align = "center"
)
library(orthoMTL)

## ----simulate-----------------------------------------------------------------
set.seed(42)

sim <- simulate_mtl(
  n = 300,
  p = 15,
  n_signals = 5,
  thresholds = c(4, 8, 14, 20),
  effect_strength = 1.2
)
sim

## ----ground-truth-------------------------------------------------------------
# True coefficients for signal features
gt <- sim$ground_truth
signal_coefs <- gt$coefficients[gt$signal_features, ]
signal_coefs

## ----ground-truth-types-------------------------------------------------------
# Effect types
gt$effect_types

## ----prepare-data-------------------------------------------------------------
thresholds <- sim$thresholds
Y <- create_longitudinal_labels(sim$SurvTime, sim$Event, thresholds)
head(Y, 10)

## ----indicator----------------------------------------------------------------
W <- create_indicator_matrix(Y)
cat("Proportion observed per threshold:\n")
colMeans(W)

## ----constraint---------------------------------------------------------------
K <- create_constraint_matrix(length(thresholds))
K

## ----fit-initial--------------------------------------------------------------
fit <- orthoMTL(
  X = sim$X, Y = Y,
  lambda = 1e-3, step_size = 0.5,
  K = K, survival = TRUE, censored.mat = W,
  alpha = 0.5
)
summary(fit)

## ----heatmap-initial, fig.width=8, fig.height=6-------------------------------
plot_heatmap(fit)

## ----cv, results='hide'-------------------------------------------------------
n_tasks <- length(sim[["thresholds"]])
folds <- rep(1:5, length.out = nrow(sim[["X"]]))

cv_res <- cv_orthoMTL(
  X.train   = sim[["X"]],
  Y.train   = Y,
  W.train   = W,
  K         = K,
  lambdas   = c(1e-5, 1e-4, 1e-3),
  alphas    = c(0, 0.5, 1),
  stepsizes = c(1, 2, 5),
  diag_vals = c(0.5, n_tasks, 2 * n_tasks),
  survival  = TRUE,
  folds     = folds,
  n_cores   = 1,
  seed      = 42,
  verbose   = FALSE
)

## ----cv-results---------------------------------------------------------------
print(cv_res)

## ----fit-final----------------------------------------------------------------
best <- cv_res$best
K_final <- K
diag(K_final) <- best$diag_val

fit_final <- orthoMTL(
  X = sim$X, Y = Y,
  lambda    = best$lambda,
  alpha     = best$alpha,
  step_size = best$stepsize,
  K         = K_final,
  survival  = TRUE,
  censored.mat = W
)
summary(fit_final)

## ----heatmap-final, fig.width=8, fig.height=6---------------------------------
plot_heatmap(fit_final)

## ----correlation, fig.width=6, fig.height=5-----------------------------------
plot_correlation(fit_final)

## ----bootstrap, results='hide'------------------------------------------------
boot_res <- bootstrap_orthoMTL(
  X = sim$X, Y = Y,
  lambda    = best$lambda,
  alpha     = best$alpha,
  step_size = best$stepsize,
  K         = K_final,
  survival  = TRUE,
  censored.mat = W,
  n_repeats = 200,
  n_cores   = 1,
  verbose   = FALSE
)

## ----bootstrap-print----------------------------------------------------------
print(boot_res)

## ----bootstrap-select---------------------------------------------------------
signal_feats <- gt[["signal_features"]]
null_feats <- gt[["null_features"]]
effect_types <- gt[["effect_types"]]

selected <- c(
  signal_feats[effect_types == "switch"][1],
  signal_feats[effect_types == "constant"][1],
  signal_feats[effect_types == "early"][1],
  null_feats[1]
)

cat("Selected features:\n")
cat("  Switch signal:  ", selected[1], "(Cox should miss)\n")
cat("  Constant signal:", selected[2], "(Cox should find)\n")
cat("  Early signal:   ", selected[3], "(Cox may dilute)\n")
cat("  Null feature:   ", selected[4], "(neither should find)\n")

## ----bootstrap-plot, fig.width=8, fig.height=6--------------------------------
suppressWarnings(
  plot_bootstrap(boot_res, features = selected)
)

## ----cox-comparison, eval=requireNamespace("glmnet", quietly = TRUE) && requireNamespace("survival", quietly = TRUE)----
library(survival)
library(glmnet)

surv_obj <- Surv(time = sim[["SurvTime"]], event = sim[["Event"]])

# Cross-validate alpha (mixing parameter)
alphas <- seq(0, 1, by = 0.1)
cv_scores <- numeric(length(alphas))

for (i in seq_along(alphas)) {
  set.seed(42)
  cvfit <- cv.glmnet(
    x = sim[["X"]],
    y = surv_obj,
    family = "cox",
    type.measure = "C",
    alpha = alphas[i]
  )
  cv_scores[i] <- cvfit[["cvm"]][cvfit[["index"]]["min", ]]
}

best_alpha <- alphas[which.max(cv_scores)]
cat("Best alpha:", best_alpha, "(CV C-index:", max(cv_scores), ")\n")

# Refit with best alpha
set.seed(42)
cox_fit <- cv.glmnet(
  x = sim[["X"]],
  y = surv_obj,
  family = "cox",
  type.measure = "C",
  alpha = best_alpha
)

cox_coefs <- as.numeric(coef(cox_fit, s = "lambda.min"))
names(cox_coefs) <- colnames(sim[["X"]])

## ----cox-table, eval=requireNamespace("glmnet", quietly = TRUE) && requireNamespace("survival", quietly = TRUE)----
# Build comparison table
ortho_mean_abs <- apply(abs(coef(fit_final)), 1, mean)

comparison <- data.frame(
  feature = gt$signal_features,
  effect_type = as.character(gt$effect_types),
  orthoMTL_mean_abs = round(ortho_mean_abs[gt$signal_features], 4),
  cox_coef = round(cox_coefs[gt$signal_features], 4),
  cox_detected = cox_coefs[gt$signal_features] != 0,
  stringsAsFactors = FALSE
)

cat("Signal feature detection comparison:\n\n")
print(comparison, row.names = FALSE)

## ----cox-vs-ortho, fig.width=8, fig.height=4, eval=requireNamespace("glmnet", quietly = TRUE) && requireNamespace("survival", quietly = TRUE)----
# Pick the switch feature — most interesting temporal pattern
switch_feat <- signal_feats[effect_types == "switch"][1]

# What Cox sees: one number
cat(switch_feat, "— Cox coefficient:", 
    round(cox_coefs[switch_feat], 3), "\n")

# What orthoMTL sees: a trajectory
cat(switch_feat, "— orthoMTL coefficients:\n")
print(round(coef(fit_final)[switch_feat, ], 3))

# What the truth is
cat(switch_feat, "— True coefficients (sign-aligned):\n")
print(round(-gt[["coefficients"]][switch_feat, ], 3))

## ----ground-truth-recovery----------------------------------------------------
true_coefs <- gt$coefficients[gt$signal_features, ]
est_coefs <- coef(fit_final)[gt$signal_features, ]

# Sign convention:
#   Simulation uses log-hazard scale (negative = protective, reduces hazard)
#   orthoMTL models P(progression-free) (positive = protective)
# Negate true coefficients to align
true_aligned <- -true_coefs

task_cors <- sapply(seq_len(ncol(true_aligned)), function(k) {
  cor(true_aligned[, k], est_coefs[, k])
})
names(task_cors) <- colnames(true_aligned)

cat("Correlation between true and estimated coefficients per threshold:\n")
print(round(task_cors, 3))

## ----recovery-heatmap, fig.width=8, fig.height=5------------------------------
make_long <- function(mat, source_label) {
  data.frame(
    feature = rep(rownames(mat), ncol(mat)),
    threshold = rep(colnames(mat), each = nrow(mat)),
    weight = as.vector(mat),
    source = source_label,
    stringsAsFactors = FALSE
  )
}

# Normalize each matrix to [-1, 1] by dividing by its own max absolute value
normalize <- function(mat) mat / max(abs(mat), na.rm = TRUE)

combined <- rbind(
  make_long(normalize(true_aligned), "True (sign-aligned)"),
  make_long(normalize(est_coefs),    "Estimated")
)

combined$feature <- factor(combined$feature, levels = rev(rownames(true_aligned)))
combined$threshold <- factor(combined$threshold, levels = colnames(true_aligned))
combined$source <- factor(combined$source, levels = c("True (sign-aligned)", "Estimated"))



# Now both panels use the full color range
ggplot2::ggplot(combined, ggplot2::aes(threshold, feature, fill = weight)) +
  ggplot2::geom_tile() +
  ggplot2::scale_fill_gradient2(low = "red", mid = "white", high = "blue", midpoint = 0,
                       limits = c(-1, 1)) +
  ggplot2::facet_wrap(~ source) +
  ggplot2::labs(x = "Threshold (months)", y = NULL, fill = "Normalized\nCoefficient") +
  ggplot2::theme_minimal()