library(tidyverse)
library(here)

theta_mean <- c(-5.490935,
                -0.0367022,
                0.768536,
                -1.344474,
                0.3740968,
                -4.0944849,
                -4.0944849,
                -3.2869827,
                log(5294^2 / sqrt(1487 ^ 2 + 5294 ^ 2)),
                1.753857,
                1.1100567,
                -0.8513930)
theta_var <- matrix(0, 12, 12)
theta_var[1:5,1:5] <- c(+4.32191e-2, -7.83e-4,    -7.247e-3,   -6.42e-4,    -5.691e-3,
                        -7.83e-4,    +2.71566e-5, +3.3e-5,     -1.1e-4,     +2.8e-8,
                        -7.247e-3,   +3.3e-5,     +1.18954e-2, +1.84e-4,    +5.1e-6,
                        -6.42e-4,    -1.11e-4,    +1.84e-4,    +1.46369e-1, +2.59e-4,
                        -5.691e-3,   +2.8e-8,     +5.1e-6,     +2.59e-4,    +2.25151e-3)
theta_var[6,6] <- 0.6527744^2
theta_var[7,7] <- 0.6527744^2
theta_var[8,8] <- 0.4889407^2
theta_var[9,9] <- log(1487^2/5294^2 + 1)
theta_var[10,10] <- 0.235360^2
theta_var[11,11] <- 0.2147719^2
theta_var[12,12] <- 0.1433964^2

state_labels  <- c("Standard_SuccessP", "Standard_Revision", "Standard_SuccessR", "Standard_Death",
                  "NP1_SuccessP", "NP1_Revision", "NP1_SuccessR", "NP1_Death")
cycle_labels  <- 1:60
output_labels <- c("Standard_Costs", "Standard_QALYs", "NP1_Costs", "NP1_QALYs", "Incremental_Costs", "Incremental_QALYs")
param_labels  <- c("beta_0", "beta_age", "beta_male", "beta_NP1", "ln_gamma", "logit_phi", "logit_psi", "logit_rho", "ln_C", "logit_U_SuccessP", "logit_U_SuccessR", "logit_U_Revision")

names(theta_mean) <- param_labels
dimnames(theta_var) <- list(theta_i = param_labels, theta_j = param_labels)

age <- 60
male <- 0

mortality <- function(age, male) {
  if (male) {
    if      (age < 45) 0.00151
    else if (age < 55) 0.00393
    else if (age < 65) 0.0109
    else if (age < 75) 0.0316
    else if (age < 85) 0.0801
    else               0.1879
  } else {
    
    if      (age < 45) 0.00099
    else if (age < 55) 0.0026
    else if (age < 65) 0.0067
    else if (age < 75) 0.0193
    else if (age < 85) 0.0535
    else               0.1548
  }
}

logit <- function(p) log(p) - log(1 - p)
expit <- function(x) 1 / (1 + exp(-x))

x    <- array(0,        dim = c(8, 60),            dimnames = list(state = state_labels, cycle = cycle_labels))
P    <- array(0,        dim = c(8, 8, 60),         dimnames = list(state_from = state_labels, state_to = state_labels, cycle = cycle_labels))
V    <- array(0,        dim = c(8, 6, 60),         dimnames = list(state = state_labels, output = output_labels, cycle = cycle_labels))
gt   <- array(NA_real_, dim = c(6, 60),            dimnames = list(output = output_labels, cycle = cycle_labels))
dx   <- array(0,        dim = c(8, 60, 12),        dimnames = list(state = state_labels, cycle = cycle_labels, theta_i = param_labels))
d2x  <- array(0,        dim = c(8, 60, 12, 12),    dimnames = list(state = state_labels, cycle = cycle_labels, theta_i = param_labels, theta_j = param_labels))
dP   <- array(0,        dim = c(8, 8, 60, 12),     dimnames = list(state_from = state_labels, state_to = state_labels, cycle = cycle_labels, theta_i = param_labels))
d2P  <- array(0,        dim = c(8, 8, 60, 12, 12), dimnames = list(state_from = state_labels, state_to = state_labels, cycle = cycle_labels, theta_i = param_labels, theta_j = param_labels))
dV   <- array(0,        dim = c(8, 6, 60, 12),     dimnames = list(state = state_labels, output = output_labels, cycle = cycle_labels, theta_i = param_labels))
d2V  <- array(0,        dim = c(8, 6, 60, 12, 12), dimnames = list(state = state_labels, output = output_labels, cycle = cycle_labels, theta_i = param_labels, theta_j = param_labels))
dgt  <- array(NA_real_, dim = c(6, 60, 12),        dimnames = list(output = output_labels, cycle = cycle_labels, theta_i = param_labels))
d2gt <- array(NA_real_, dim = c(6, 60, 12, 12),    dimnames = list(output = output_labels, cycle = cycle_labels, theta_i = param_labels, theta_j = param_labels))

# INITIAL STATE DISTRIBUTION
x[4, 1] <- expit(theta_mean[6])
x[1, 1] <- 1 - x[4, 1]
x[8, 1] <- expit(theta_mean[6])
x[5, 1] <- 1 - x[8, 1]

# TRANSITION MATRIX
P[4, 4, ] <- 1
P[8, 8, ] <- 1
for (t in 1:60) {
  rt_Standard <- 1 - exp(exp(theta_mean[1] + theta_mean[2] * age + theta_mean[3] * male) * ((t - 1)^exp(theta_mean[5]) - t^exp(theta_mean[5])))
  rt_NP1      <- 1 - exp(exp(theta_mean[1] + theta_mean[2] * age + theta_mean[3] * male + theta_mean[4]) * ((t - 1)^exp(theta_mean[5]) - t^exp(theta_mean[5])))
  P[1, 2, t]  <- rt_Standard
  P[5, 6, t]  <- rt_NP1
  P[1, 4, t]  <- mortality(age + t, male)
  P[2, 4, t]  <- mortality(age + t, male) + expit(theta_mean[7])
  P[3, 4, t]  <- mortality(age + t, male)
  P[5, 8, t]  <- mortality(age + t, male)
  P[6, 8, t]  <- mortality(age + t, male) + expit(theta_mean[7])
  P[7, 8, t]  <- mortality(age + t, male)
  P[3, 2, t]  <- expit(theta_mean[8])
  P[7, 6, t]  <- expit(theta_mean[8])
  P[1, 1, t]  <- 1 - P[1, 2, t] - P[1, 4, t]
  P[2, 3, t]  <- 1 - P[2, 4, t]
  P[3, 3, t]  <- 1 - P[3, 2, t] - P[3, 4, t]
  P[5, 5, t]  <- 1 - P[5, 6, t] - P[5, 8, t]
  P[6, 7, t]  <- 1 - P[6, 8, t]
  P[7, 7, t]  <- 1 - P[7, 6, t] - P[7, 8, t]
}

# VALUE MATRIX
V_base <- matrix(0, 8, 6)
V_base[2, 1] <- exp(theta_mean[9])
V_base[2, 5] <- -exp(theta_mean[9])
V_base[6, c(3, 5)] <- exp(theta_mean[9])

V_base[1, 2] <- expit(theta_mean[10])
V_base[1, 6] <- -expit(theta_mean[10])
V_base[5, 4] <- expit(theta_mean[10])
V_base[5, 6] <- expit(theta_mean[10])

V_base[3, 2] <- expit(theta_mean[11])
V_base[3, 6] <- -expit(theta_mean[11])
V_base[7, 4] <- expit(theta_mean[11])
V_base[7, 6] <- expit(theta_mean[11])

V_base[2, 2] <- expit(theta_mean[12])
V_base[2, 6] <- -expit(theta_mean[12])
V_base[6, 4] <- expit(theta_mean[12])
V_base[6, 6] <- expit(theta_mean[12])

for (t in 0:59) {
  V[, , t+1] <- V_base %*% diag(rep(c(1.06^(-t), 1.015^(-t)), 3), 6, 6)
}

# STANDARD MARKOV COHORT SIMULATION
g0 <- c(394, 0, 579, 0, 579 - 394, 0)
gt[, 1] <- t(V[, , 1]) %*% x[, 1]
for (t in 2:60) {
  x[, t]  <- t(P[, , t-1]) %*% x[, t-1]
  gt[, t] <- t(V[, , t]) %*% x[, t]
}
g <- g0 + rowSums(gt)

# DELTA METHOD
# - Transition matrix (use symbolic differentiation, but exclude mortality as constant)
P11 <- expression(exp(exp(theta1 + theta2 * age + theta3 * male) * ((t - 1)^exp(theta5) - t^exp(theta5))))
P12 <- expression(1 - exp(exp(theta1 + theta2 * age + theta3 * male) * ((t - 1)^exp(theta5) - t^exp(theta5))))
P23 <- expression(1 - 1/(1+exp(-theta7)))
P24 <- expression(1/(1+exp(-theta7)))
P32 <- expression(1/(1+exp(-theta8)))
P33 <- expression(1 - 1/(1+exp(-theta8)))
P55 <- expression(exp(exp(theta1 + theta2 * age + theta3 * male + theta4) * ((t - 1)^exp(theta5) - t^exp(theta5))))
P56 <- expression(1 - exp(exp(theta1 + theta2 * age + theta3 * male + theta4) * ((t - 1)^exp(theta5) - t^exp(theta5))))
P67 <- expression(1 - 1/(1+exp(-theta7)))
P68 <- expression(1/(1+exp(-theta7)))
P76 <- expression(1/(1+exp(-theta8)))
P77 <- expression(1 - 1/(1+exp(-theta8)))

dP11 <- deriv3(P11, paste0("theta", 1:12))
dP12 <- deriv3(P12, paste0("theta", 1:12))
dP23 <- deriv3(P23, paste0("theta", 1:12))
dP24 <- deriv3(P24, paste0("theta", 1:12))
dP32 <- deriv3(P32, paste0("theta", 1:12))
dP33 <- deriv3(P33, paste0("theta", 1:12))
dP55 <- deriv3(P55, paste0("theta", 1:12))
dP56 <- deriv3(P56, paste0("theta", 1:12))
dP67 <- deriv3(P67, paste0("theta", 1:12))
dP68 <- deriv3(P68, paste0("theta", 1:12))
dP76 <- deriv3(P76, paste0("theta", 1:12))
dP77 <- deriv3(P77, paste0("theta", 1:12))

param_list <- set_names(as.list(theta_mean), paste0("theta", 1:12))
grad <- purrr::attr_getter("gradient")
hessian <- purrr::attr_getter("hessian")

from_ <- rep(c(1, 2, 3, 5, 6, 7), each = 2)
to_   <- c(1:4, 2, 3, 5:8, 6, 7)

for (t in 1:60) {
  param_list$t <- pmax(t, 1 + 1e-10)
  
  walk2(from_, to_, function(f_, t_) {
    tmp <- eval(get(paste0("dP", f_, t_)), param_list)
    dP[f_, t_, t, ] <<- unname(grad(tmp))[1, ]
    d2P[f_, t_, t, , ] <<- unname(hessian(tmp))[1, , ]
  })
}

# - Value matrix
#   - Utilities
dV[1, 2, , 10] <- -exp(-theta_mean[10]) * (1 + exp(-theta_mean[10])) ^ (-2) * 1.015 ^ (-(1:60)) # V[1, 2, t] = U_SuccessP 1.015^(-t)
dV[1, 6, , 10] <- -dV[1, 2, , 10]
dV[2, 2, , 12] <- -exp(-theta_mean[12]) * (1 + exp(-theta_mean[12])) ^ (-2) * 1.015 ^ (-(1:60))
dV[2, 6, , 12] <- -dV[2, 2, , 12]
dV[3, 2, , 11] <- -exp(-theta_mean[11]) * (1 + exp(-theta_mean[11])) ^ (-2) * 1.015 ^ (-(1:60))
dV[3, 6, , 11] <- -dV[3, 2, , 11]
dV[5, 4, , 10] <- -exp(-theta_mean[10]) * (1 + exp(-theta_mean[10])) ^ (-2) * 1.015 ^ (-(1:60))
dV[5, 6, , 10] <- dV[5, 4, , 10]
dV[6, 4, , 12] <- -exp(-theta_mean[12]) * (1 + exp(-theta_mean[12])) ^ (-2) * 1.015 ^ (-(1:60))
dV[6, 6, , 12] <- dV[6, 4, , 12]
dV[7, 4, , 11] <- -exp(-theta_mean[11]) * (1 + exp(-theta_mean[11])) ^ (-2) * 1.015 ^ (-(1:60))
dV[7, 6, , 11] <- dV[7, 4, , 11]

d2V[1, 2, , 10, 10] <- exp(-theta_mean[10]) * (1 - exp(-theta_mean[10])) * (1 + exp(-theta_mean[10])) ^ (-3) * 1.015 ^ (-(1:60))
d2V[1, 6, , 10, 10] <- -d2V[1, 2, , 10, 10]
d2V[2, 2, , 12, 12] <- exp(-theta_mean[12]) * (1 - exp(-theta_mean[12])) * (1 + exp(-theta_mean[12])) ^ (-3) * 1.015 ^ (-(1:60))
d2V[2, 6, , 12, 12] <- -d2V[2, 2, , 12, 12]
d2V[3, 2, , 11, 11] <- exp(-theta_mean[11]) * (1 - exp(-theta_mean[11])) * (1 + exp(-theta_mean[11])) ^ (-3) * 1.015 ^ (-(1:60))
d2V[3, 6, , 11, 11] <- -d2V[3, 2, , 11, 11]
d2V[5, 4, , 10, 10] <- exp(-theta_mean[10]) * (1 - exp(-theta_mean[10])) * (1 + exp(-theta_mean[10])) ^ (-3) * 1.015 ^ (-(1:60))
d2V[5, 6, , 10, 10] <- d2V[5, 4, , 10, 10]
d2V[6, 4, , 12, 12] <- exp(-theta_mean[12]) * (1 - exp(-theta_mean[12])) * (1 + exp(-theta_mean[12])) ^ (-3) * 1.015 ^ (-(1:60))
d2V[6, 6, , 12, 12] <- d2V[6, 4, , 12, 12]
d2V[7, 4, , 11, 11] <- exp(-theta_mean[11]) * (1 - exp(-theta_mean[11])) * (1 + exp(-theta_mean[11])) ^ (-3) * 1.015 ^ (-(1:60))
d2V[7, 6, , 11, 11] <- d2V[7, 4, , 11, 11]

#   - Costs
dV[2, 1, , 9] <- V[2, 1, ] # V[2, 1, t] = exp(theta[9]) 1.06^(-t)
dV[2, 5, , 9] <- V[2, 5, ]
dV[6, 3, , 9] <- V[6, 3, ]
dV[6, 5, , 9] <- V[6, 5, ]
d2V[2, 1, , 9, 9] <- V[2, 1, ]
d2V[2, 5, , 9, 9] <- V[2, 5, ]
d2V[6, 3, , 9, 9] <- V[6, 3, ]
d2V[6, 5, , 9, 9] <- V[6, 5, ]

# - Initial state distribution
dx[c(1, 4, 5, 8), 1, 6] <- c(-1, 1, -1, 1) * exp(-theta_mean[6]) * (1 + exp(-theta_mean[6])) ^ (-2)
d2x[c(1, 4, 5, 8), 1, 6, 6] <- c(1, -1, 1, -1) * exp(-theta_mean[6]) * (1 - exp(-theta_mean[6])) * (1 + exp(-theta_mean[6])) ^ (-3)
for (i in 1:12) {
  dgt[, 1, i] <- t(dV[, , 1, i]) %*% x[, 1] + t(V[, , 1]) %*% dx[, 1, i]
  for (j in 1:12) {
    d2gt[, 1, i, j] <- t(d2V[, , 1, i, j]) %*% x[, 1] +
      t(dV[, , 1, i]) %*% dx[, 1, j] +
      t(dV[, , 1, j]) %*% dx[, 1, i] +
      t(V[, , 1]) %*% d2x[, 1, i, j]
  }
}

# - Recurrence relation
for (t in 2:60) {
  for (i in 1:12) {
    dx[, t, i] <- t(dP[, , t-1, i]) %*% x[, t-1] + t(P[, , t-1]) %*% dx[, t-1, i]
    dgt[, t, i] <- t(dV[, , t, i]) %*% x[, t] + t(V[, , t]) %*% dx[, t-1, i]
    for (j in 1:12) {
      d2x[, t, i, j] <- t(d2P[, , t-1, i, j]) %*% x[, t-1] +
        t(dP[, , t-1, i]) %*% dx[, t-1, j] +
        t(dP[, , t-1, j]) %*% dx[, t-1, i] +
        t(P[, , t-1]) %*% d2x[, t-1, i, j]
      
      d2gt[, t, i, j] <- t(d2V[, , t, i, j]) %*% x[, t] +
        t(dV[, , t, i]) %*% dx[, t, j] +
        t(dV[, , t, j]) %*% dx[, t, i] +
        t(V[, , t]) %*% d2x[, t, i, j]
    }
  }
}

# RESULTS!
dg <- apply(dgt, c(1, 3), sum)
d2g <- apply(d2gt, c(1, 3, 4), sum)

g_expected <- sapply(1:6, function(v) g[v] + 0.5 * sum(d2g[v, , ] * theta_var))
g_variance <- dg %*% theta_var %*% t(dg)
g_sd <- sqrt(diag(g_variance))

# COMPARE WITH PSA
g_psa <- readRDS(here("THR-MCPSA.rds"))

library(patchwork)

p_standard_costs <- ggplot(g_psa, aes(x = Standard_Costs)) +
  geom_density() +
  stat_function(fun = dnorm, args = list(mean = g_expected[1], sd = g_sd[1]), linetype = "dashed", colour = "blue") +
  scale_x_continuous("Lifetime costs", labels = scales::dollar_format(prefix = "£")) +
  ggtitle("Costs for standard prosthesis") +
  theme(axis.text.y = element_blank(), axis.ticks.y = element_blank(), axis.title.y = element_blank())
p_standard_QALYs <- ggplot(g_psa, aes(x = Standard_QALYs)) +
  geom_density() +
  stat_function(fun = dnorm, args = list(mean = g_expected[2], sd = g_sd[2]), linetype = "dashed", colour = "blue") +
  scale_x_continuous("Quality adjusted life years (QALYs)") +
  ggtitle("QALYs for standard prosthesis") +
  theme(axis.text.y = element_blank(), axis.ticks.y = element_blank(), axis.title.y = element_blank())
p_NP1_costs <- ggplot(g_psa, aes(x = NP1_Costs)) +
  geom_density() +
  stat_function(fun = dnorm, args = list(mean = g_expected[3], sd = g_sd[3]), linetype = "dashed", colour = "blue") +
  scale_x_continuous("Lifetime costs", labels = scales::dollar_format(prefix = "£")) +
  ggtitle("Costs for NP1 prosthesis") +
  theme(axis.text.y = element_blank(), axis.ticks.y = element_blank(), axis.title.y = element_blank())
p_NP1_QALYs <- ggplot(g_psa, aes(x = NP1_QALYs)) +
  geom_density() +
  stat_function(fun = dnorm, args = list(mean = g_expected[4], sd = g_sd[4]), linetype = "dashed", colour = "blue") +
  scale_x_continuous("Quality adjusted life years (QALYs)") +
  ggtitle("QALYs for NP1 prosthesis") +
  theme(axis.text.y = element_blank(), axis.ticks.y = element_blank(), axis.title.y = element_blank())
p_inc_costs <- ggplot(g_psa, aes(x = Incremental_Costs)) +
  geom_density() +
  stat_function(fun = dnorm, args = list(mean = g_expected[5], sd = g_sd[5]), linetype = "dashed", colour = "blue") +
  scale_x_continuous("Lifetime costs", labels = scales::dollar_format(prefix = "£")) +
  ggtitle("Incremental costs") +
  theme(axis.text.y = element_blank(), axis.ticks.y = element_blank(), axis.title.y = element_blank())
p_inc_QALYs <- ggplot(g_psa, aes(x = Incremental_QALYs)) +
  geom_density() +
  stat_function(fun = dnorm, args = list(mean = g_expected[6], sd = g_sd[6]), linetype = "dashed", colour = "blue") +
  scale_x_continuous("Quality adjusted life years (QALYs)") +
  ggtitle("Incremental QALYs") +
  theme(axis.text.y = element_blank(), axis.ticks.y = element_blank(), axis.title.y = element_blank())

(p_standard_costs + p_NP1_costs + p_inc_costs) / (p_standard_QALYs + p_NP1_QALYs + p_inc_QALYs)
ggsave(here("Delta-PSA comparison.png"), width = 20.28, height = 10.62, units = "cm")

wtp <- seq(0, 10000, by = 50)

ceac_psa <- sapply(wtp, function(lambda) mean(lambda * g_psa$Incremental_QALYs >= g_psa$Incremental_Costs))

m <- - g_expected[5] + wtp * g_expected[6]
s <- sqrt(wtp ^ 2 * g_variance[6,6] - 2 * wtp * g_variance[5,6] + g_variance[5,5])

ggplot(tibble(x = wtp, ceac.delta = pnorm(m/s), ceac.psa = ceac_psa), aes(x)) +
  geom_line(aes(y = ceac.delta, linetype = "Delta-PSA", colour = "Delta-PSA")) +
  geom_line(aes(y = ceac.psa, linetype = "MC-PSA", colour = "MC-PSA")) +
  scale_x_continuous("Willingness-to-pay for additional QALY", labels = scales::dollar_format(prefix = "£")) +
  scale_y_continuous("Probability NP1 is cost-effective", labels = scales::percent_format()) +
  scale_colour_discrete("Method") +
  scale_linetype_discrete("Method")
ggsave(here("Delta-PSA CEAC.png"), width = 20.28, height = 10.62, units = "cm", scale = 0.8)

evpi_psa <- sapply(wtp, function(lambda) {
  inmb <- lambda * g_psa$Incremental_QALYs - g_psa$Incremental_Costs
  evpi <- mean(pmax(0, inmb)) - pmax(0, mean(inmb))
  evpi
})
evpi_delta <- m * pnorm(m/s) + s * dnorm(m/s) - pmax(0, m)

ggplot(tibble(x = wtp, evpi.delta = evpi_delta, evpi.psa = evpi_psa), aes(x)) +
  geom_line(aes(y = evpi.delta, linetype = "Delta-PSA", colour = "Delta-PSA")) +
  geom_line(aes(y = evpi.psa, linetype = "MC-PSA", colour = "MC-PSA")) +
  scale_x_continuous("Willingness-to-pay for additional QALY", labels = scales::dollar_format(prefix = "£")) +
  scale_y_continuous("Expected value of perfect information", labels = scales::dollar_format(prefix = "£")) +
  scale_colour_discrete("Method") +
  scale_linetype_discrete("Method")
ggsave(here("Delta-PSA EVPI.png"), width = 20.28, height = 10.62, units = "cm", scale = 0.8)