# ============================================================
# 03_machine_learning.R
# 三种机器学习方法筛选核心基因
# ============================================================

library(randomForest)
library(glmnet)
library(caret)
library(e1071)
library(pROC)
library(tidyverse)
library(VennDiagram)

# ---- 读取数据 ----
expr <- read.csv("data/GSE38974_COPD_expression_clean.csv", 
                 row.names = 1, check.names = FALSE)
sample_info <- read.csv("data/GSE38974_COPD_sample_info.csv")

# 假设已有关键基因集（差异基因交集）
# 实际应从差异分析结果筛选
deg_sig <- read.csv("output/COPD_DEGs_significant.csv", row.names = 1)
common_genes <- rownames(head(deg_sig, 200))  # 取Top 200示例

ml_data <- t(expr[common_genes, ])
ml_data <- as.data.frame(ml_data)
ml_data$Group <- factor(sample_info$Group[match(rownames(ml_data), sample_info$Sample)],
                        levels = c("Control", "COPD"))

# ---- 划分训练集/测试集 ----
set.seed(123)
train_idx <- createDataPartition(ml_data$Group, p = 0.7, list = FALSE)
train_data <- ml_data[train_idx, ]
test_data <- ml_data[-train_idx, ]

train_x <- train_data[, -ncol(train_data)]
train_y <- train_data$Group
test_x <- test_data[, -ncol(test_data)]
test_y <- test_data$Group

# ============================================================
# 方法1：随机森林 (Random Forest)
# ============================================================
cat("\n===== Random Forest =====\n")

set.seed(123)
rf_model <- randomForest(
  x = train_x,
  y = train_y,
  ntree = 500,
  mtry = sqrt(ncol(train_x)),
  importance = TRUE
)

rf_pred <- predict(rf_model, test_x)
rf_conf <- confusionMatrix(rf_pred, test_y)
cat("RF 准确率:", rf_conf$overall["Accuracy"], "\n")

rf_importance <- as.data.frame(importance(rf_model))
rf_importance$Gene <- rownames(rf_importance)
rf_importance <- rf_importance %>% arrange(desc(MeanDecreaseAccuracy))
rf_top50 <- head(rf_importance, 50)$Gene

# ============================================================
# 方法2：LASSO
# ============================================================
cat("\n===== LASSO =====\n")

x_train <- as.matrix(train_x)
y_train <- ifelse(train_y == "COPD", 1, 0)
x_test <- as.matrix(test_x)
y_test_num <- ifelse(test_y == "COPD", 1, 0)

set.seed(123)
cv_fit <- cv.glmnet(x = x_train, y = y_train, family = "binomial", 
                    alpha = 1, nfolds = 10, type.measure = "class")

best_lambda <- cv_fit$lambda.min
cat("最优lambda:", best_lambda, "\n")

lasso_model <- glmnet(x_train, y_train, family = "binomial", alpha = 1, lambda = best_lambda)
coef_matrix <- as.matrix(coef(lasso_model))
lasso_genes <- rownames(coef_matrix)[coef_matrix[, 1] != 0]
lasso_genes <- lasso_genes[lasso_genes != "(Intercept)"]
cat("LASSO 保留基因数:", length(lasso_genes), "\n")

lasso_prob <- predict(lasso_model, newx = x_test, type = "response")
lasso_pred <- ifelse(lasso_prob > 0.5, 1, 0)
cat("LASSO 准确率:", mean(lasso_pred == y_test_num), "\n")

# ============================================================
# 方法3：SVM-RFE
# ============================================================
cat("\n===== SVM-RFE =====\n")

rfe_data <- train_data
rfe_data$Group <- make.names(rfe_data$Group)

rfe_ctrl <- rfeControl(
  functions = caretFuncs,
  method = "cv",
  number = 5,
  verbose = FALSE
)

set.seed(123)
svmProfile <- rfe(
  x = rfe_data[, -ncol(rfe_data)],
  y = rfe_data$Group,
  sizes = c(5, 10, 20, 30, 50),
  rfeControl = rfe_ctrl,
  method = "svmRadial"
)

svm_genes <- predictors(svmProfile)
cat("SVM-RFE 选中基因数:", length(svm_genes), "\n")

# ============================================================
# 取交集
# ============================================================
cat("\n===== 交集分析 =====\n")

core_genes <- Reduce(intersect, list(rf_top50, lasso_genes, svm_genes))
cat("三种方法交集基因:", length(core_genes), "\n")

if (length(core_genes) < 3) {
  intersect_2 <- union(
    intersect(rf_top50, lasso_genes),
    union(intersect(rf_top50, svm_genes), intersect(lasso_genes, svm_genes))
  )
  core_genes <- head(intersect_2, 10)
  cat("两两交集Top 10:\n")
}

cat("核心基因:", paste(core_genes, collapse = ", "), "\n")

# 保存
write.csv(data.frame(Gene = core_genes), "output/core_genes_final.csv", row.names = FALSE)

# 韦恩图
venn <- venn.diagram(
  x = list(RF = rf_top50, LASSO = lasso_genes, SVM_RFE = svm_genes),
  filename = "output/10_ML_venn.png",
  fill = c("#E69F00", "#56B4E9", "#009E73"),
  alpha = 0.5,
  cex = 2,
  cat.cex = 1.2
)

# ROC曲线（用核心基因简化模型）
if (length(core_genes) > 0) {
  simple_train <- train_data[, c(core_genes, "Group")]
  simple_test <- test_data[, c(core_genes, "Group")]
  simple_model <- glm(Group ~ ., data = simple_train, family = binomial)
  simple_prob <- predict(simple_model, simple_test, type = "response")
  
  roc_obj <- roc(test_data$Group, simple_prob)
  png("output/11_core_genes_ROC.png", width = 800, height = 600, res = 100)
  plot(roc_obj, main = paste("ROC (AUC =", round(auc(roc_obj), 3), ")"))
  dev.off()
}

cat("\n机器学习分析完成！\n")