# ============================================================
# 教程配套脚本 08：生存分析与预后模型
# 输入：表达矩阵（TCGA-LIHC）+ 临床生存数据
# 输出：KM 曲线、森林图、LASSO 模型、ROC
# ============================================================

library(survival)
library(survminer)
library(glmnet)
library(timeROC)
library(tidyverse)

# ---- 1. 准备 TCGA 生存数据（示例结构；正式数据来自 Xena phenotype 文件）----
# 假设已读取：
#   expr_tcga: 表达矩阵（基因 × 样本，样本名如 TCGA-XX-XXXX-01A）
#   clin: 数据框，列 = c("sample", "OS.time", "OS", "age", "gender", "stage")
# 教学演示：构造一个示例数据框跑通流程
set.seed(42)
n <- 371
clin <- data.frame(
  sample = paste0("TCGA-", sprintf("%04d", 1:n)),
  OS.time = round(rexp(n, rate = 1/30) + 1, 1),   # 模拟生存月数
  OS = rbinom(n, 1, 0.45),                        # 模拟结局 0 生存 1 死亡
  age = round(runif(n, 30, 85)),
  gender = sample(c("MALE", "FEMALE"), n, replace = TRUE),
  stage = sample(c("I", "II", "III", "IV"), n, replace = TRUE,
                 prob = c(0.4, 0.3, 0.2, 0.1))
)
head(clin)

# 模拟 6 个焦亡基因的表达（演示用随机数；正式用真实表达矩阵）
gene_names <- c("GSDMD", "GSDME", "CASP1", "IL1B", "NLRP3", "AIM2")
expr_demo <- matrix(rnorm(n * length(gene_names), mean = 6, sd = 1.5),
                    nrow = length(gene_names), ncol = n,
                    dimnames = list(gene_names, clin$sample))

# ---- 2. 单基因生存分析（以 GSDMD 为例）----
gene <- "GSDMD"
expr_g <- expr_demo[gene, ]

# 中位数分组：高表达 vs 低表达
group_g <- ifelse(expr_g > median(expr_g), "High", "Low")
clin$group <- group_g

fit_km <- survfit(Surv(OS.time, OS) ~ group, data = clin)
p_km <- ggsurvplot(fit_km,
                   data = clin,
                   pval = TRUE,                     # 显示 log-rank P
                   risk.table = TRUE,
                   palette = c("#C44E52", "#4C72B0"),
                   legend.labs = c("High", "Low"),
                   xlab = "Time (months)", ylab = "Overall survival",
                   title = paste0("KM curve: ", gene))
ggsave("figures/R_km_curve.png", print(p_km), width = 8, height = 6, dpi = 300)

# ---- 3. 单因素 / 多因素 Cox 回归 ----
# 把基因表达加入临床协变量
clin_cox <- clin %>%
  mutate(expr = expr_g)

cox_uni <- coxph(Surv(OS.time, OS) ~ expr, data = clin_cox)
summary(cox_uni)

cox_multi <- coxph(Surv(OS.time, OS) ~ expr + age + gender + stage, data = clin_cox)
summary(cox_multi)

# 森林图（多因素）
p_forest <- ggforest(cox_multi, data = clin_cox)
ggsave("figures/R_forest_plot.png", p_forest, width = 8, height = 5, dpi = 300)

# ---- 4. LASSO 构建多基因风险模型 ----
# 用演示的 6 个基因（正式分析用全部焦亡基因或候选 hub 基因）
X <- t(expr_demo)                       # 样本 × 基因
y <- Surv(clin$OS.time, clin$OS)

set.seed(123)
cv_fit <- cv.glmnet(X, y, family = "cox", alpha = 1, nfolds = 10)
plot(cv_fit)                            # 交叉验证曲线
best_lambda <- cv_fit$lambda.min

# 提取 LASSO 选中的基因与系数
lasso_coef <- coef(cv_fit, s = "lambda.min")
selected <- lasso_coef@Dimnames[[1]][lasso_coef@i + 1]
coefs <- lasso_coef@x
cat("LASSO 选中的基因:\n"); print(data.frame(gene = selected, coef = coefs))

# 风险评分 = Σ(系数 × 表达)
risk_score <- as.numeric(X %*% as.matrix(lasso_coef))
clin$risk <- risk_score

# ---- 5. 时间依赖 ROC（1/3/5 年）----
roc_res <- timeROC(T = clin$OS.time, delta = clin$OS,
                   marker = clin$risk, cause = 1,
                   times = c(12, 36, 60), iid = TRUE)
roc_res$AUC
# 画 ROC
plot(roc_res, time = 12, col = "#C44E52", title = FALSE)
plot(roc_res, time = 36, col = "#DD8452", add = TRUE)
plot(roc_res, time = 60, col = "#4C72B0", add = TRUE)
legend("bottomright", legend = c("1-year", "3-year", "5-year"),
       col = c("#C44E52", "#DD8452", "#4C72B0"), lwd = 2)
savePlot("figures/R_roc.png", type = "png")   # Windows 可用；跨平台用 pdf/png 设备

cat("生存分析与预后模型完成！\n")
