rm(list = ls())
library(glmnet)
library(survival)

set.seed(2025)

# Fonte dos dados:
# hastie.su.domains/StatLearnSparsity_files/DATA/lymphoma.html

# Dados do estudo de linfoma
data_x <- read.table('./lymphx.txt', header = FALSE)
data_y <- read.table('./lymphtim.txt', header = FALSE)
data_cens <- read.table('./lymphstatus.txt', header = FALSE)

dim(data_x)
dim(data_y)
dim(data_cens)

head(data_y)
head(data_cens)
table(data_cens)

# matriz de preditores
x <- as.matrix(data_x)
# objeto Surv com tempos de sobrevivencia e indicador de censura
y <- survival::Surv(data_y$V1, data_cens$V1)
n <- dim(x)[1]
p <- dim(x)[2]

# Estimativa da função de sobrevivência (Kaplan-Meier)
km_fit <- survival::survfit(y ~ 1)

# png(file="linfoma_kaplan.png", 
#     width=600, height=500, res = 100)
plot(km_fit, xlab = 'Tempos', ylab = 'S(t)', conf.int = 0.95)
# dev.off()

###################################
# ajuste do lasso

lasso_fit = glmnet(x, y, family = "cox", alpha = 1)

# png(file="linfoma_lasso_path.png", 
#     width=600, height=500, res = 100)
plot(lasso_fit, xvar = 'lambda', sign.lambda = 1, 
     lwd=2, ylab='Coeficientes')
# dev.off()

# validação cruzada para escolher lambda
cv_lasso_fit <- glmnet::cv.glmnet(x, y, family = 'cox', alpha = 1, 
                                  type.measure = 'deviance', nfolds = 10) 

# png(file="linfoma_lasso_deviance_cv.png", 
#     width=600, height=500, res = 100)
plot(cv_lasso_fit, sign.lambda = 1)
# dev.off()

# coeficientes estimados usando o lambda 'ótimo'
cv_lasso_fit$lambda.min
id_labda_min <- which(lasso_fit$lambda == cv_lasso_fit$lambda.min)
beta_hat_lasso <- coef(lasso_fit, s=lasso_fit$lambda[id_labda_min])

sum(beta_hat_lasso != 0)

# png(file="linfoma_estimativas_lasso.png", 
#     width=600, height=500, res = 100)
par(mar = c(5,4,1,1), mgp = c(2.5, 1, 0))
plot(1:p, rep(0, p), ylim = c(min(beta_hat_lasso), max(beta_hat_lasso)), 
     pch=' ', ylab = expression(beta[i]), xlab = 'i', cex.lab=1.5)
id <- which(beta_hat_lasso == 0)
points(id, beta_hat_lasso[id], pch=20, col='red', cex=2)
id <- which(beta_hat_lasso != 0)
points(id, beta_hat_lasso[id], pch=8, col='blue', cex=2)
legend("topright", c(expression(hat(beta[i])==0), expression(hat(beta[i])!=0)),
       pch = c(20, 8), col = c("red", "blue"), pt.cex=1.5)
# dev.off()

# valores ajustados do preditor linear
fitted_eta_values <- x%*%beta_hat_lasso # modelo de Cox não tem beta_0

# Kaplan-Meier para casos com o preditor linear ajustado positivo
id_etapos <- which(fitted_eta_values > 0)
km_fit_eta_pos <- survival::survfit(y[id_etapos] ~ 1)

# Kaplan-Meier para casos com o preditor linear ajustado negativo
id_etaneg <- which(fitted_eta_values <= 0)
km_fit_eta_neg <- survival::survfit(y[id_etaneg] ~ 1)


# png(file="linfoma_kaplan_etavalues.png", 
#    width=600, height=500, res = 100)
plot(km_fit, xlab = 'Tempos', ylab = 'S(t)', 
     col = 'Black', lwd=2, conf.int = FALSE)
lines(km_fit_eta_pos, conf.int = FALSE, col = 'Red', lwd=2)
lines(km_fit_eta_neg, conf.int = FALSE, col = 'Green', lwd=2)
legend("topright", c(expression(hat(eta)<=0), expression(hat(eta)>0), 'todos'),
       fill = c('Green', 'Red', 'Black'), cex=1.2)
# dev.off()





