# Load libraries
library(MASS)
library(dplyr)
library(ggplot2)
library(lubridate)

# Load data (assumes file already read into 'data')
data <- read.csv("PedITS_Data.csv")
data$reg_date <- as.Date(data$reg_date)

# Create time variables
data <- data %>%
  arrange(reg_date) %>%
  mutate(
    time = as.numeric(reg_date - min(reg_date)),
    intervention = ifelse(reg_date >= as.Date("2020-03-01"), 1, 0),
    time_post = ifelse(intervention == 1, time - time[which(reg_date == as.Date("2020-03-01"))] + 1, 0),
    t = 1:n(),
    log_pop = log(pop_all)
  )

# Add Fourier terms (2 harmonics for seasonality)
K <- 2
for (k in 1:K) {
  data[[paste0("sin", k)]] <- sin(2 * pi * k * data$t / 12)
  data[[paste0("cos", k)]] <- cos(2 * pi * k * data$t / 12)
}

# Fit Negative Binomial model
model <- glm.nb(ed_urg ~ time + intervention + time_post +
                  sin1 + cos1 + sin2 + cos2 +
                  offset(log_pop),
                data = data)

# Predict fitted values
data$fitted <- predict(model, type = "response")

# Create counterfactual dataset (no intervention)
cf_data <- data
cf_data$intervention <- 0
cf_data$time_post <- 0
cf_data$counterfactual <- predict(model, newdata = cf_data, type = "response")

# ---- Extract Slopes and Intercepts ----

# Slope (log scale)
slope_pre <- coef(model)["time"]
slope_post <- coef(model)["time"] + coef(model)["time_post"]

# Percent change per month
percent_pre <- (exp(slope_pre) - 1) * 100
percent_post <- (exp(slope_post) - 1) * 100

# Intercepts at March 2020
intervention_date <- as.Date("2020-03-01")
i <- which(data$reg_date == intervention_date)
X_actual <- c(1, data$time[i], 1, 0, as.numeric(data[i, c("sin1","cos1","sin2","cos2")]), data$log_pop[i])
X_cf     <- c(1, data$time[i], 0, 0, as.numeric(data[i, c("sin1","cos1","sin2","cos2")]), data$log_pop[i])
beta     <- coef(model)
log_actual <- sum(beta[1:8] * X_actual[1:8]) + X_actual[9]
log_cf     <- sum(beta[1:8] * X_cf[1:8]) + X_cf[9]
intercept_actual <- exp(log_actual)
intercept_cf     <- exp(log_cf)

# Final values
final_idx <- nrow(data)
final_actual <- data$fitted[final_idx]
final_cf <- cf_data$counterfactual[final_idx]

# ---- Print summary ----
cat("Pre-intervention slope (%/month):", round(percent_pre, 4), "%\n")
cat("Post-intervention slope (%/month):", round(percent_post, 4), "%\n")
cat("Intercept at March 2020 - Actual:", round(intercept_actual), "\n")
cat("Intercept at March 2020 - Counterfactual:", round(intercept_cf), "\n")
cat("Final predicted count - Actual:", round(final_actual), "\n")
cat("Final predicted count - Counterfactual:", round(final_cf), "\n")

# ---- Plot ----
pre_data <- data[data$reg_date < intervention_date, ]
post_data <- data[data$reg_date >= intervention_date, ]
cf_post_data <- cf_data[cf_data$reg_date >= intervention_date, ]

ggplot(data, aes(x = reg_date)) +
  geom_line(aes(y = ed_urg), color = "black", size = 0.4) +  # actual
  geom_line(data = pre_data, aes(y = fitted), color = "red", size = 1) +  # pre-trend
  geom_line(data = post_data, aes(y = fitted), color = "red", size = 1) +  # post-trend
  geom_line(data = cf_post_data, aes(y = counterfactual), color = "blue", linetype = "dashed", size = 1) +  # counterfactual
  geom_vline(xintercept = intervention_date, color = "green", linetype = "dashed", size = 1) +
  labs(title = "A", x = "Time", y = "ED visits per month") +
  theme_minimal(base_size = 14) +
  theme(
    plot.title = element_text(hjust = 0, face = "bold", size = 16),
    axis.title.x = element_text(margin = margin(t = 10)),
    axis.title.y = element_text(margin = margin(r = 10))
  )

coef_intervention <- coef(model)["intervention"]
se_intervention <- summary(model)$coefficients["intervention", "Std. Error"]

rr <- exp(coef_intervention)
ci_lower <- exp(coef_intervention - 1.96 * se_intervention)
ci_upper <- exp(coef_intervention + 1.96 * se_intervention)

cat("Rate Ratio (RR):", round(rr, 4), "\n")
cat("95% CI:", round(ci_lower, 4), "-", round(ci_upper, 4), "\n")

summary(model)

# Reshape data: urgent vs. nonurgent into one column
data_long <- data %>%
  select(reg_date, pop_all, ed_urg, ed_nonurg) %>%
  pivot_longer(
    cols = c(ed_urg, ed_nonurg),
    names_to = "visit_type",
    values_to = "visits"
  ) %>%
  arrange(reg_date) %>%
  mutate(
    time = as.numeric(reg_date - min(reg_date)),
    intervention = ifelse(reg_date >= as.Date("2020-03-01"), 1, 0),
    time_post = ifelse(intervention == 1,
                       time - time[which(reg_date == as.Date("2020-03-01"))] + 1,
                       0),
    t = rep(1:(nrow(data)), times = 2), # time index repeats for each type
    log_pop = log(pop_all),
    visit_type = factor(visit_type, levels = c("ed_urg", "ed_urg"))
  )

# Add Fourier terms (seasonality)
K <- 2
for (k in 1:K) {
  data_long[[paste0("sin", k)]] <- sin(2 * pi * k * data_long$t / 12)
  data_long[[paste0("cos", k)]] <- cos(2 * pi * k * data_long$t / 12)
}

model_compare <- glm.nb(
  visits ~ time * visit_type +
    intervention * visit_type +
    time_post * visit_type +
    sin1 + cos1 + sin2 + cos2 +
    offset(log_pop),
  data = data_long
)

summary(model_compare)
exp(coef(model_compare)) # Relative risks for each term