---
title: "Linear ITR: Application"
author: "Andong Wang"
date: "4/4/2023"
output: html_document
---

```{r setup, include=FALSE}
knitr::opts_chunk$set(echo = TRUE)
library(tidyverse)
library(rgenoud)
library(foreach)
library(doParallel)
library(reshape2)
library(ggridges)
# Import cleaned and wrangled data
mimic_base <- read_csv("./cleandata/mimic_base.csv")
eicu_base <- read_csv("./cleandata/eicu_base.csv")
# Required functions to run beforehand
source("./TLITR Functions.R")
```

# Finding calibration weights

```{r getCW}
set.seed(2023)

rct <- mimic_base
rwd <- eicu_base

# remove outliers
rct <- rct %>% filter(Temp_C<50 & Temp_C>25)

# covariates matrices for both population
x_rwd <- as.matrix(rwd[,c("readmit","meanbp","temperature","respiratoryrate", 
                          "sodium","glucose","bun","creatinine","bilirubin",
                          "albumin","wbc", "admissionweight","age"
                          )])
x_rct <- as.matrix(rct[,c("re_admission","MeanBP","Temp_C", "RR",
                      "Sodium", "Glucose","BUN", "Creatinine","Total_bili",
                      "Albumin","WBC_count","Weight_kg","age"
                      )])

# Obtain calibration weights
mu_g0 <- apply(x_rwd, MARGIN = 2, FUN = mean)

lamFun <- function(lam, moments, moments.bar){ ## vector lam and x
  qi <- (exp(moments%*%lam)/sum(exp(moments%*%lam)))[,1]
  colSums(qi*moments) - moments.bar
}

res = nleqslv::searchZeros(matrix(rnorm(length(mu_g0)*10,0, 0.02),nrow = 10),
                                lamFun, moments = x_rct, moments.bar = mu_g0)
lam.hat <- nleqslv::searchZeros(matrix(rnorm(length(mu_g0)*10,0, 0.02),nrow = 10),
                                lamFun, moments = x_rct, moments.bar = mu_g0)$x[1,]

wts <- exp(x_rct%*%lam.hat)/sum(exp(x_rct%*%lam.hat))
```

# Density Plots to show heterogeneity

```{r ridge}
df_x_rct <- x_rct %>%
  as.data.frame() %>%
  mutate(Population="Source: MIMIC-III")
colnames(df_x_rct) <- c("Re-admission", "Mean Blood Pressure", "Body Temperature", "Respiratory Rate", "Sodium", "Glucose", "Blood Urea Nitrogen", "Creatinine", "Bilirubin", "Albumin", "WBC Count", "Weight", "Age", "Population")

df_x_rwd <- x_rwd %>%
  as.data.frame() %>%
  mutate(Population="Target: eICU-CRD")
colnames(df_x_rwd) <- c("Re-admission", "Mean Blood Pressure", "Body Temperature", "Respiratory Rate", "Sodium", "Glucose", "Blood Urea Nitrogen", "Creatinine", "Bilirubin", "Albumin", "WBC Count", "Weight", "Age", "Population")

df_ridge <- rbind(df_x_rct, df_x_rwd) %>%
  select(-"Re-admission")

long_data <- melt(df_ridge, id.vars = "Population", variable.name = "Feature", value.name = "Value")
long_data$Population <- factor(long_data$Population, levels = c("Target: eICU-CRD", "Source: MIMIC-III"))

p3 <- ggplot(long_data, aes(x = Value, color = Population, fill = Population)) +
  geom_density(alpha = 0.5) +
  facet_wrap(~ Feature, scales = "free") +  # Create a panel for each feature
  labs(#title = "Density Plots for Numerical Covariates",
       x = "Value",
       y = "Density") +
  theme_minimal() +
  theme(legend.position = "bottom")
p3
```

# Covariates distribution comparison: weighted source population vs target population

```{r weighted-density}
# Selected covariates
selected_vars <- c("Mean Blood Pressure", "Body Temperature", "Respiratory Rate","Albumin")

# Target population
df_target <- x_rwd %>%
  as.data.frame() %>%
  mutate(Population = "Target: eICU-CRD",
         wts = 1)
colnames(df_target) <- c("Re-admission", "Mean Blood Pressure", "Body Temperature", "Respiratory Rate",
                         "Sodium", "Glucose", "Blood Urea Nitrogen", "Creatinine", "Bilirubin", 
                         "Albumin", "WBC Count", "Weight", "Age", "Population", "wts")

# Unweighted source
df_source_unweighted <- x_rct %>%
  as.data.frame() %>%
  mutate(Population = "Source: MIMIC-III (Unweighted)",
         wts = 1)
colnames(df_source_unweighted) <- c("Re-admission", "Mean Blood Pressure", "Body Temperature", "Respiratory Rate",
                                    "Sodium", "Glucose", "Blood Urea Nitrogen", "Creatinine", "Bilirubin", 
                                    "Albumin", "WBC Count", "Weight", "Age", "Population", "wts")

# Weighted source
df_source_weighted <- x_rct %>%
  as.data.frame() %>%
  mutate(Population = "Source: MIMIC-III (Weighted)",
         wts = wts)
colnames(df_source_weighted) <- c("Re-admission", "Mean Blood Pressure", "Body Temperature", "Respiratory Rate",
                                  "Sodium", "Glucose", "Blood Urea Nitrogen", "Creatinine", "Bilirubin", 
                                  "Albumin", "WBC Count", "Weight", "Age", "Population", "wts")

# Combine all three
df_all <- bind_rows(df_target, df_source_unweighted, df_source_weighted) %>%
  select(all_of(selected_vars), Population, wts)

# Reshape to long format
long_all <- melt(df_all, id.vars = c("Population", "wts"),
                 variable.name = "Feature", value.name = "Value")

# Set custom legend order
long_all$Population <- factor(long_all$Population,
                              levels = c("Target: eICU-CRD",
                                         "Source: MIMIC-III (Unweighted)",
                                         "Source: MIMIC-III (Weighted)"))

# Plot
p3.5 <- ggplot(long_all, aes(x = Value, weight = wts, color = Population,fill = Population)) +
  geom_density(alpha = 0.5, adjust = 1.2) +
  facet_wrap(~ Feature, scales = "free") +
  labs(
    #title = "Covariate Distribution Comparison After Calibration",
    x = "Value",
    y = "Density"
  ) +
  scale_color_manual(values = c(
    "Target: eICU-CRD" = "#F8766D",              # Red
    "Source: MIMIC-III (Unweighted)" = "#00BFC4",# Blue
    "Source: MIMIC-III (Weighted)" = "#7CAE00"   # Green
  )) +
  scale_fill_manual(values = c(
    "Target: eICU-CRD" = "#F8766D",
    "Source: MIMIC-III (Unweighted)" = "#00BFC4",
    "Source: MIMIC-III (Weighted)" = "#7CAE00"
  )) +
  theme_minimal(base_size = 14) +
  theme(legend.position = "bottom")

p3.5
```

# Estimate values for propensity score model component of the value function

```{r prop-score}
# toc stands for treatment, outcome, and covariates
rct_toc <- as.data.frame(cbind(
  treatment = rct$mechvent, outcome = 1-rct$death, x_rct)) # outcome is survival rate, so we can maximize it
prop_model <- glm(formula = treatment ~ .-outcome, data = rct_toc, family = binomial)
summary(prop_model)
prop_scores <- prop_model$fitted.values # checking propensity score
rct_toc <- rct_toc %>% mutate(PS = prop_scores)
pred_prop <- predict(prop_model, newdata = rct_toc, type="response")
```

# Estimate values for outcome regression model component of the value function

```{r or-est}
OR_model1 <-
  lm(formula = outcome ~ .-treatment,
     data = rct_toc %>% filter(treatment == 1))
summary(OR_model1)

OR_model0 <-
  lm(formula = outcome ~ .-treatment,
     data = rct_toc %>% filter(treatment == 0))
summary(OR_model0)

#### Predict using lm
pred_OR1 <- predict(OR_model1, newdata = rct_toc)
pred_OR0 <- predict(OR_model0, newdata = rct_toc)
rct_toc <- cbind(rct_toc, pred_OR1, pred_OR0)
OR_est <- mean(rct_toc$pred_OR1) - mean(rct_toc$pred_OR0)
```

# Define the value function 

```{r value-func}
value_weighted_AIPW <- function(eta) {
  reg <- (Xd%*%eta>0)*1
  Qd <- pred_OR1 * reg + pred_OR0 * (1-reg)
  Cd <- reg == Assignment
  pid <- pred_prop * reg + p0 * (1-reg)
  value <- Cd * Y / pid - (Cd - pid) / pid * Qd
  return(sum(weights*value))
}
```

# Define global variables for value function 

```{r def-variables}
Xd <- cbind(intercept = 1, x_rct)
Assignment <- rct_toc$treatment
Y <- rct_toc$outcome
p0 <- 1-pred_prop
```


# Multiple iterations of GA optimization for unweighted ITR

```{r uwga}
# Customizable: Number of iterations of GA
n_iter <- 300

# Set up parameters and variables for data collection
wts0 <- rep(1/nrow(x_rct), nrow(x_rct))
weights <- wts0
value_uw <- rep(0,n_iter)
itr_uw <- matrix(0, nrow = n_iter, ncol = n_var)

start_time <- Sys.time()
for (i in 1:n_iter) {
  seed <- 100+i
  set.seed(seed)
  print(paste0("Iteration: ", i, "; Seed: ", seed))
  dom <- matrix(c(rep(-1,n_var),rep(1,n_var)),n_var,2)
  result_opt <- genoud(fn=value_weighted_AIPW,nvars=n_var,max=TRUE,pop.size=3000,
                       starting.values=rep(0,n_var),Domains=dom,optim.method="Nelder-Mead")
  itr_uw[i,] <- result_opt$par 
  value_uw[i] <- result_opt$value
}
end_time <- Sys.time()
end_time - start_time
```

# Evaluate unweighted ITRs and select one with the highest value

```{r eval-itrs}
weights <- wts
for (i in 1:n_iter) {
  value_uw[i] <- value_weighted_AIPW(itr_uw[i,])
}
iter_uw_opt <- which(value_uw == max(value_uw)); iter_uw_opt
itr_uw_opt <- itr_uw[iter_uw_opt,]; itr_uw_opt
value_uw_opt <- value_uw[iter_uw_opt[1]]; value_uw_opt
# seed 48 produces the highest valued unweighted ITR out of seed 1~100: 0.7220347
# seed 202 produces the highest valued unweighted ITR out of seed 101~400: 0.7222514
```

# Multiple iterations of GA optimization for weighted ITR

```{r wga}
# Customizable: Number of iterations of GA
n_iter <- 100

# Set up parameters and variables for data collection
weights <- wts
n_var <- ncol(x_rct)+1
value_w <- rep(0,n_iter)
itr_w <- matrix(0, nrow = n_iter, ncol = n_var)

start_time <- Sys.time()
for (i in 1:n_iter) {
  seed <- i
  set.seed(seed)
  print(paste0("Iteration: ", i, "; Seed: ", seed))
  dom <- matrix(c(rep(-1,n_var),rep(1,n_var)),n_var,2)
  result_opt <- genoud(fn=value_weighted_AIPW,nvars=n_var,max=TRUE,pop.size=3000,
                       starting.values=rep(0,n_var),Domains=dom,optim.method="Nelder-Mead")
  itr_w[i,] <- result_opt$par
  value_w[i] <- result_opt$value
}
end_time <- Sys.time()
time_diff_w <- end_time - start_time
```

# Inspect values and select the weighted ITR with the highest value

```{r inspect}
iter_w_opt <- which(value_w == max(value_w)); 
itr_w_opt <- itr_w[iter_w_opt,]; itr_w_opt
value_w_opt <- value_w[iter_w_opt[1]]; value_w_opt
```

# 300 iterations of GA optimization for weighted ITR

```{r 300}
# Customizable: the seed number that 300 iterations start with
starting_seed <- 700

# Set up parameters and variables for data collection
weights <- wts
n_var <- ncol(x_rct)+1
n_iter <- 300
value300 <- 0
itr300 <- rep(0, n_var)
seed_opt <- 0

start_time <- Sys.time()
for (i in 1:n_iter) {
  seed <- starting_seed+i
  set.seed(seed)
  print(paste0("Iteration: ", i, "; Seed: ", seed))
  dom <- matrix(c(rep(-1,n_var),rep(1,n_var)),n_var,2)
  result_opt <- genoud(fn=value_weighted_AIPW,nvars=n_var,max=TRUE,pop.size=3000,
                       starting.values=rep(0,n_var),Domains=dom,optim.method="Nelder-Mead")
  if (result_opt$value > value300) {
    value300 <- result_opt$value
    itr300 <- result_opt$par
    seed_opt <- seed
  }
}
end_time <- Sys.time()
end_time-start_time
# 300 iterations takes roughly 30 minutes to run on my personal laptop.

# seed 15 produces the highest valued weighted ITR out of seed 1~100: 0.7595422
# seed 114 produces the highest valued weighted ITR out of seed 101~400: 0.7595471
# seed 588 produces the highest valued weighted ITR out of seed 401~700: 0.7612141
# seed 866 produces the highest valued weighted ITR out of seed 701~1000: 0.7615876
```

# Single Iteration of GA optimization for weighted ITR

```{r single}
# Customizable: Seed number for the single iteration of GA
seed <- 866

# Set up parameters and variables for data collection
weights <- wts
set.seed(seed)
print(paste0("Iteration: ", i, "; Seed: ", seed))
dom <- matrix(c(rep(-1, n_var), rep(1, n_var)), n_var, 2)
result_opt <-
  genoud(
    fn = value_weighted_AIPW,
    nvars = n_var,
    max = TRUE,
    pop.size = 3000,
    starting.values = rep(0, n_var),
    Domains = dom,
    optim.method = "Nelder-Mead"
  )

itr_w_opt <- result_opt$par
value_w_opt <- result_opt$value
```

# Rescale covariates to adjust coefficients, visualize coefficients, select important covariates based on the magnitude of coefficients of the optimal ITR.

```{r coef-adj}
x_rct_sd <- apply(x_rct, 2, sd)
coef_adj <- itr_w_opt * c(1,x_rct_sd)
var_label <- c("Intercept", "Re-admission", "Mean Blood Pressure", "Body Temperature", "Respiratory Rate", "Sodium", "Glucose", "Blood Urea Nitrogen", "Creatinine", "Bilirubin", "Albumin", "WBC Count", "Weight", "Age")
df_plot <- data.frame(Covariate = var_label, AdjustedCoefficient = coef_adj, Coefficient = itr_w_opt)
df_plot <- df_plot %>% 
  mutate(Sign = as.factor(sign(AdjustedCoefficient)))
df_plot$AdjustedCoefficient <- abs(df_plot$AdjustedCoefficient)
df_plot <- df_plot %>% 
  arrange(desc(AdjustedCoefficient)) 
  # filter(Coefficient > 1)
p4 <- ggplot(data=df_plot, aes(x = reorder(Covariate, -AdjustedCoefficient), 
                         y = AdjustedCoefficient, 
                         fill = Sign)) +
  geom_bar(stat = "identity") +
  geom_text(aes(label = round(AdjustedCoefficient,2)), vjust=-0.5) +
  scale_y_continuous(breaks = seq(0,30,5),limits = c(0,30)) +
  # ggtitle("Compare Covariates' Coefficients in the Optimal ITR") +
  xlab("Patient Covariate") +
  ylab("Adjusted Coefficient Estimate") +
  theme(
    axis.text.x = element_text(angle = 90)
  )
p4
```

# Optimal ITR result

$$

-0.39327497\cdot Glucose + 0.65068207\cdot Blood Urea Nitrogen + 0.62822781\cdot Age  -0.24842725\cdot Weight + 0.43325392\cdot Mean Blood Pressure -0.47377776\cdot WBC Count + 0.87998515\cdot Respiratory Rate + 0.88298164\cdot Bilirubin  -0.62201393\cdot Sodium  -0.05645087\cdot Creatinine -0.76444845\cdot Re-admission + 0.55447353\cdot Body Temperature + 0.36333257\cdot Albumin + 0.16342131 > 0

$$

# Estimate variance of value function estimates

```{r}

```

