library(strucchange)
library(dplyr)
library(tidyr)
library(ggplot2)
library(ggpubr)
library(here)

# Simulate BP data
set.seed(20)
n <- 70
switch_time <- 50
monitor_time = 20

hr <- 25 + rpois(n, lambda = 5)
baseline_bp <- rnorm(n, mean=72 - seq(n) * 0.05 * (seq(n) > switch_time), sd=3)

x <- data.frame(
  hr = hr,
  baseline_bp = baseline_bp
)

init_beta <- matrix(c(0.1,1), nrow = 2, ncol = 1)
new_beta <- matrix(c(0.15,0.9), nrow = 2, ncol = 1)
eps <- rnorm(n, mean = 0, sd=1)
init_future_bp <- as.matrix(x[seq(switch_time),]) %*% init_beta + eps[seq(switch_time)]
new_future_bp <- as.matrix(x[seq(switch_time + 1, n),]) %*% new_beta + eps[seq(switch_time + 1, n)]

all_dat <- data.frame(
  outcome_bp = c(init_future_bp, new_future_bp),
  hr = hr,
  baseline_bp = baseline_bp
)

# Monitor changes in the X1
monitor_dat <- all_dat[seq(monitor_time),]
me.x1 <- mefp(hr ~ 1, type="OLS-CUSUM", rescale=TRUE,
                data=monitor_dat, alpha=0.05)
monitor_dat <- all_dat
me.x1 <- monitor(me.x1)
print(me.x1)
plot(me.x1)

# Monitor changes in the X2
monitor_dat <- all_dat[seq(monitor_time),]
me.x2 <- mefp(baseline_bp ~ 1, type="OLS-CUSUM", rescale=TRUE,
              data=monitor_dat, alpha=0.05)
monitor_dat <- all_dat
me.x2 <- monitor(me.x2)
print(me.x2)
plot(me.x2)

# Monitor changes in the Y | X
monitor_dat <- all_dat[seq(monitor_time),]
me.mefp <- mefp(outcome_bp ~ ., type="OLS-CUSUM", rescale=TRUE,
                data=monitor_dat, alpha=0.05)

monitor_dat <- all_dat[seq(n),]
me.mefp <- monitor(me.mefp)
print(me.mefp)
plot(me.mefp, functional = NULL)

# Aggregate into final plot
monitor_len <- length(me.x1$process)
l_bounds_plot_df <- data.frame(
  time = seq(monitor_len),
  var_hr_cusum = -me.x1$border(seq(monitor_time + 1, n)),
  var_baseline_bp_cusum = -me.x2$border(seq(monitor_time + 1, n)),
  var_model_cusum = -me.mefp$border(seq(monitor_time + 1, n))
)
u_bounds_plot_df <- data.frame(
  time = seq(monitor_len),
  var_hr_cusum = me.x1$border(seq(monitor_time + 1, n)),
  var_baseline_bp_cusum = me.x2$border(seq(monitor_time + 1, n)),
  var_model_cusum = me.mefp$border(seq(monitor_time + 1, n))
)

stat_plot_df <- data.frame(
  time = seq(length(me.x1$process)),
  var_hr_cusum = me.x1$process,
  var_baseline_bp_cusum = me.x2$process,
  var_model_cusum = me.mefp$process
)
stat_plot_df <- pivot_longer(stat_plot_df, cols = starts_with("var"), names_prefix="var_")
l_bounds_plot_df <- pivot_longer(l_bounds_plot_df, cols = starts_with("var"), names_prefix="var_")
u_bounds_plot_df <- pivot_longer(u_bounds_plot_df, cols = starts_with("var"), names_prefix="var_")
stat_plot_df$label = "CUSUM"
l_bounds_plot_df$label = "CUSUM lower bound"
u_bounds_plot_df$label = "CUSUM upper bound"

plot_df <- rbind(stat_plot_df, l_bounds_plot_df, u_bounds_plot_df)

mod_future_map <- lm(outcome_bp ~ ., data = all_dat[seq(monitor_time),])
pred_future_map <- predict(mod_future_map, 
                           newdata = all_dat[-seq(monitor_time),c(2,3)])
MSE <- sqrt((pred_future_map - all_dat$outcome_bp[-seq(monitor_time)])^2)
MSE_plot_df <- data.frame(time = seq(n-monitor_time), 
                          name = rep("MSE", n-monitor_time), 
                          value = MSE, 
                          label = rep("MSE", n-monitor_time))
plot_df <- rbind(plot_df, MSE_plot_df)

name.labs <- c("Average Baseline HR Monitoring: Fluctuation Process with Boundaries", "Average Baseline MAP Monitoring: Fluctuation Process with Boundaries", "MAP Prediction Model Monitoring: Fluctuation Process with Boundaries", "MAP Prediction Model Monitoring: RMSE")
names(name.labs) <- c("hr_cusum", "baseline_bp_cusum", "model_cusum", "MSE")


pdf(file = here("manuscript", "strucchange_plot.pdf"), width = 12, height = 8)
ggplot(plot_df, aes(x = time, y = value, color = label)) +  
  geom_line(size = 1.5) + 
  scale_color_manual(values = c("#1F78B4", "#A6CEE3", "#A6CEE3", "#7570B3")) +
  facet_wrap(~name, labeller = labeller(name = name.labs), ncol=1, 
             scales = "free_y") + 
  theme_bw() +
  labs(x = "Time", title = "") + 
  ylab(NULL) +
  theme(legend.title = element_blank(), legend.position = "none") +
  theme(text = element_text(size = 20)) +
  theme(panel.spacing.y=unit(1, "lines"))
dev.off()