library(tidyverse)

set.seed(13470)
N_PSA <- 20000

revision_risk_coef <- c(-5.490935, -0.0367022, +0.768536, -1.344474, +0.3740968)
revision_risk_var  <- matrix(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),
                             5, 5)
revision_risk_psa  <- MASS::mvrnorm(N_PSA, mu = revision_risk_coef, Sigma = revision_risk_var)

params <- tibble(
  beta_0     = revision_risk_psa[, 1],
  beta_age   = revision_risk_psa[, 2],
  beta_male  = revision_risk_psa[, 3],
  beta_NP1   = revision_risk_psa[, 4],
  ln_gamma   = revision_risk_psa[, 5],
  phi        = rbeta(N_PSA, 2, 98),
  psi        = rbeta(N_PSA, 2, 98),
  rho        = rbeta(N_PSA, 4, 96),
  C          = rgamma(N_PSA, shape = 12.6749407, scale = 417.6745372),
  U_SuccessP = rbeta(N_PSA, 119.566667, 21.1),
  U_SuccessR = rbeta(N_PSA, 87.140625, 29.046875),
  U_Revision = rbeta(N_PSA, 69.7, 162.633333)
)

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")

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))

runPSA <- function(...) {
  theta <- flatten_dbl(rlang::list2(...))
  
  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))
  
  # INITIAL STATE DISTRIBUTION
  x[4, 1] <- theta[6]
  x[1, 1] <- 1 - x[4, 1]
  x[8, 1] <- theta[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[1] + theta[2] * age + theta[3] * male) * ((t - 1)^exp(theta[5]) - t^exp(theta[5])))
    rt_NP1      <- 1 - exp(exp(theta[1] + theta[2] * age + theta[3] * male + theta[4]) * ((t - 1)^exp(theta[5]) - t^exp(theta[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) + theta[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) + theta[7]
    P[7, 8, t]  <- mortality(age + t, male)
    P[3, 2, t]  <- theta[8]
    P[7, 6, t]  <- theta[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] <- theta[9]
  V_base[2, 5] <- -theta[9]
  V_base[6, c(3, 5)] <- theta[9]
  
  V_base[1, 2] <- theta[10]
  V_base[1, 6] <- -theta[10]
  V_base[5, 4] <- theta[10]
  V_base[5, 6] <- theta[10]
  
  V_base[3, 2] <- theta[11]
  V_base[3, 6] <- -theta[11]
  V_base[7, 4] <- theta[11]
  V_base[7, 6] <- theta[11]
  
  V_base[2, 2] <- theta[12]
  V_base[2, 6] <- -theta[12]
  V_base[6, 4] <- theta[12]
  V_base[6, 6] <- theta[12]
  
  for (t in 1:60) {
    V[, , t] <- 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)
  
  # RESULTS
  as.list(g)
}

g_psa <- pmap_dfr(params, runPSA)

saveRDS(g_psa, "THR-MCPSA.rds")