#---------------------------------------------------------
# Scriptname:  3.two_stage_country_specific.R
#
# Purpose of script: Analyse the simulated case-crossover
# dataset by country and meta-analyze coefficients (Fig 5)
#
#
#  ! Analyses are run on a simulated dataset.       !
#  ! Results and Figure will differ from results    ! 
#  ! presented in the manuscript.                   !
#  ! In this script we assume that NUTS "1" and "2" !
#  ! belongs to country "A", while NUTS "3" and "4" !
#  ! to country "B".                                !
#
# Purpose of script:
#
# Analyse the simulated case-crossover dataset by country 
# and meta-analyze coefficients (Fig 5)
#         
#  ------------------------------------------------------
#  ! This script is run on simulated data:              !
#  !                                                    !
#  ! - "case_cross_data.csv" is the simulated           !
#  !    case-crossover dataset created with script      !
#  !    "1.from_cases_to_case_crossover.R"              !
#  !                                                    !
#  ! In this script we assume that NUTS "1" and "2"     !
#  ! belongs to country "A", while NUTS "3" and "4"     !
#  ! to country "B".                                    !
#  !                                                    !
#  ! Since this script is run on a simulated datase     !
#  ! results and figure obtained will slightly differ   !
#  ! from those presented in the manuscript.            !
#  ------------------------------------------------------
#
# Author(s): Giovenale Moirano
# Date Created: 05 May 2024
# Email: giovenale.moirano@bsc.es
#
#---------------------------------------------------------
# 1. Clean Environment
#---------------------------------------------------------

rm(list = ls())

#------------------------------------------------------------
# 2.Set WD
#------------------------------------------------------------

# setwd("mydirectory")

#---------------------------------------------------------
# 3. Loadpackages
#---------------------------------------------------------

packages <- c( 
  "tidyr", "lubridate", "dplyr",
  "mixmeta", "survival", "dlnm"
)
install.packages(setdiff(packages, rownames(installed.packages())))
lapply(packages, require, character.only = TRUE)


#------------------------------------------------------------
# 4a. Load CC_dataset
#------------------------------------------------------------

data <- read.csv("data/7.case_cross_data.csv")
ftable(data$caco, data$NUTS_ID)


# Define Strata for conditional logistic regression
data$clogit_id <- as.numeric(substr(data$ID, 3, 6))


#------------------------------------------------------------
# 4b. Nuts 1 and 2 belongs to country A; Nuts 3 and 4 belongs
# to country B
#------------------------------------------------------------

data$country[data$NUTS_ID %in% c("1", "2")] <- "A"
data$country[data$NUTS_ID %in% c("3", "4")] <- "B"


#------------------------------------------------------------
# 5. Create lagged exposures Matrix for Weekly Average Temp
# in the 8 weeks preceding the index date
#------------------------------------------------------------

# 5a. For each case and control period, extract daily temperatures
# based on the starting prefix ("tg")

t_mat <- as.matrix(data[, which(substr(names(data), 1, 2) == "tg")])
dim(t_mat)

# 5b. Compute weekly average temperature

a <- NULL
for (i in seq(from = 7, to = 56, by = 7)) {
  a[[ceiling(i / 7)]] <- apply(t_mat[, (i - 6):i], 1, mean)
}

t_mat_avg <- as.matrix(as.data.frame(a))
hist(t_mat_avg)

#------------------------------------------------------------
# 6. Create lagged exposures matrix for Weekly Cumulative Prep
# in the  8 weeks preceding the index date
#------------------------------------------------------------

# 6a. Extract daily precipitation based on the prefix ("pr")

p_mat <- as.matrix(data[, which(substr(names(data), 1, 2) == "pr")])
dim(p_mat)


# 6b. Compute weekly cumulative prep

a2 <- NULL
for (i in seq(from = 7, to = 56, by = 7)) {
  a2[[ceiling(i / 7)]] <- apply(p_mat[, (i - 6):i], 1, sum)
}

p_mat_avg <- as.matrix(as.data.frame(a2))
dim(p_mat_avg)

hist(p_mat_avg)

#------------------------------------------------------------
# 7. Combine Data with Exposure data
#------------------------------------------------------------

data <- data %>% dplyr::select(
  country, NUTS_ID, id,
  caco, sampling_date
)

t_mat_avg <- as.data.frame(t_mat_avg)
names(t_mat_avg) <- c(paste0("tg_lag_", 1:8))
dim(t_mat_avg)

p_mat_avg <- as.data.frame(p_mat_avg)
names(p_mat_avg) <- c(paste0("p_lag_", 1:8))
dim(p_mat_avg)

multi_country_mat <- cbind(data, t_mat_avg, p_mat_avg)
multi_country_mat <- as.data.frame(multi_country_mat)
names(multi_country_mat)


#------------------------------------------------------------
# 8. Load time series of two meteorological parameters and
# compute percentiles for country A and country B.
#------------------------------------------------------------

clima <- read.csv(file = "data/6.weekly_prep_temp.csv")
clima <- clima %>%
  mutate(
    week = lubridate::week(date),
    year = lubridate::year(date),
    country = ifelse(NUTS_ID %in% c("1", "2"), "A", "B")
  ) %>%
  filter(week >= 19 & week <= 44) %>%
  group_by(week, year, country) %>%
  summarise(t_mean = mean(tg, na.rm = T), prep_cum = sum(prep, na.rm = T)) %>%
  group_by(substr(country, 1, 2)) %>%
  summarise(
    t_25 = quantile(t_mean, probs = 0.25, na.rm = TRUE),
    t_33 = quantile(t_mean, probs = 0.33, na.rm = TRUE),
    t_50 = quantile(t_mean, probs = 0.50, na.rm = TRUE),
    t_66 = quantile(t_mean, probs = 0.66, na.rm = TRUE),
    t_75 = quantile(t_mean, probs = 0.75, na.rm = TRUE),
    t_90 = quantile(t_mean, probs = 0.90, na.rm = TRUE),
    p_25 = quantile(prep_cum, probs = 0.25, na.rm = TRUE),
    t_33 = quantile(t_mean, probs = 0.33, na.rm = TRUE),
    p_50 = quantile(prep_cum, probs = 0.50, na.rm = TRUE),
    t_66 = quantile(t_mean, probs = 0.66, na.rm = TRUE),
    p_75 = quantile(prep_cum, probs = 0.75, na.rm = TRUE),
    p_90 = quantile(prep_cum, probs = 0.90, na.rm = TRUE)
  )

clima

#------------------------------------------------------------
# 9. Organise the dataset as a list of dataframe based on
# the country name
#------------------------------------------------------------

country <- as.character(unique(multi_country_mat$country))
datalist <- lapply(country, function(country) multi_country_mat[multi_country_mat$country == country, ])
names(datalist) <- country

#------------------------------------------------------------
# 9. Create empty matrices to store the coef and vcov for each
# meterological factor
#------------------------------------------------------------

coef_t <- matrix(NA, length(datalist), 6, dimnames = list(country, paste0("b", seq(6))))
S_t <- vector("list", length(datalist))
coef_p <- matrix(NA, length(datalist), 6, dimnames = list(country, paste0("b", seq(6))))
S_p <- vector("list", length(datalist))
names(S_t) <- names(S_p) <- country

#------------------------------------------------------------
# 9. Define parameters for the cross-basis of Temp and Prep
# and run analyses looping per country and using country-specific
# percentiles. 
#------------------------------------------------------------

for (i in seq(datalist)) {
  # specify arguments of the cross-basis for temp and prep based on
  # country specific percentiles
  tasknots_t <- unlist(c(clima[[i, "t_33"]], clima[[i, "t_66"]]))
  argvar_t <- list(fun = "ns", knots = tasknots_t, intercept = F)
  argvar_p <- list(fun = "lin", intercept = F)
  lag_knots <- c(1.66, 4.32)
  arglag_t <- list(fun = "bs", knots = lag_knots, intercept = T)
  cb_t <- crossbasis(datalist[[i]][, which(substr(names(datalist[[i]]), 1, 2) == "tg")], argvar = argvar_t, arglag = arglag_t)
  arglag_p <- list(fun = "bs", knots = lag_knots, intercept = T)
  cb_p <- crossbasis(datalist[[i]][, which(substr(names(datalist[[i]]), 1, 2) == "p_")], argvar = argvar_p, arglag = arglag_p)
  # run country specific model
  m <- clogit(caco ~ cb_t + cb_p + strata(clogit_id),
    method = "approximate",
    data = datalist[[i]]
  )
  # crossreduce the basis to obtain lag-specific estimates
  cr_t <- crossreduce(cb_t, m, type = "var", by = 0.5, value = round(unlist(clima[[i, "t_75"]])) + 1, cen = round(unlist(clima[[i, "t_75"]])))
  cr_p <- crossreduce(cb_p, m, type = "var", value = round(unlist(clima[i, "p_75"])) + 10, cen = round(unlist(clima[i, "p_75"])))
  # store country specific coefficients and vcov
  coef_t[i, ] <- coef(cr_t)
  S_t[[i]] <- vcov(cr_t)
  coef_p[i, ] <- coef(cr_p)
  S_p[[i]] <- vcov(cr_p)
  print(i)
}


#------------------------------------------------------------
# 9. Fit multivariate meta-analyses for temperature and plot results
#------------------------------------------------------------

# Fit meta-analyses
mix_t <- mixmeta(coef_t ~ 1, S_t, method = "ml")
print(summary(mix_t), digits = 3)

# Estimate blups
mixblup_t <- mixmeta::blup(mix_t, vcov = TRUE)

# Recreate spline transformation over the lag dim
lag <- seq(0, 7, by = 1)
lagvar <- list(fun = "bs", knots = lag_knots, intercept = T)
blag_t <- do.call("onebasis", c(list(x = lag), lagvar))

# Pooled estimates
predpool_t <- crosspred(blag_t,
  coef = coef(mix_t), vcov = vcov(mix_t),
  model.link = "log"
)

# Country-specific estimates
predblup_t <- lapply(seq(nrow(coef_t)), function(i) {
  crosspred(blag_t,
    coef = mixblup_t[[i]]$blup,
    vcov = mixblup_t[[i]]$vcov,
    model.link = "log"
  )
})

# Plot results
plot(predpool_t,
  type = "l", ylab = "OR", ylim = c(0.80, 1.4), lwd = 2,
  xlab = "Lags (week)", main = "A",
  xaxt = "n", ci = "area", col = "red"
)
custom_x_ticks <- c(0:7)
custom_x_labels <- c(paste(1:8))
axis(1, at = custom_x_ticks, labels = custom_x_labels)

for (i in 1:length(predblup_t)) {
  lines(predblup_t[[i]],
    type = "l", ylab = "OR", ylim = c(0.9, 1.3), lwd = 2,
    xlab = "Lags (week)",
    # main = country[i],
    col = rgb(255, 0, 0, alpha = 60, maxColorValue = 255),
    xaxt = "n"
  )
}

#------------------------------------------------------------
# 10. Fit multivariate meta-analyses for prep and plot results
#------------------------------------------------------------

# Fit meta-analyses
mix_p <- mixmeta(coef_p ~ 1, S_p, method = "ml")
print(summary(mix_p), digits = 3)

# Estimate blups
mixblup_p <- mixmeta::blup(mix_p, vcov = TRUE)

# Recreate spline transformation over the lag dim
blag <- do.call("onebasis", c(list(x = lag), lagvar))

# Pooled Estimates 
predpool_p <- crosspred(blag, coef = coef(mix_p), vcov = vcov(mix_p), model.link = "log")

# Country-specific Estimates 
predblup_p <- lapply(seq(nrow(coef_p)), function(i) {
  crosspred(blag,
    coef = mixblup_p[[i]]$blup,
    vcov = mixblup_p[[i]]$vcov,
    model.link = "log"
  )
})

plot(predpool_p,
  type = "l", ci = "area", ylab = "OR", ylim = c(0.95, 1.21), lwd = 2,
  xlab = "Lags (week)", main = "B", xaxt = "n", col = "forestgreen"
)
custom_x_ticks <- c(0:7)
custom_x_labels <- c(paste(1:8))
axis(1, at = custom_x_ticks, labels = custom_x_labels)

for (i in 1:length(predblup_p)) {
  lines(predblup_p[[i]],
    ylab = "OR", ylim = c(0.9, 1.3), lwd = 2,
    xlab = "Lags (week)",
    # main = country[i],
    col = rgb(34, 139, 34, alpha = 60, maxColorValue = 255),
    xaxt = "n", pch = 18, type = "l"
  )
}

