#------------------------------------------------------------------------------#
####              Calculate cost-effectiveness outcomes                     ####
#------------------------------------------------------------------------------#
#' Calculate cost-effectiveness outcomes
#'
#' \code{calculate_ce_out} calculates costs and effects for a given vector of 
#' parameters using a simulation model.
#' @param l_params_all List with all parameters of decision model
#' @param n_wtp Willingness-to-pay threshold to compute net benefits
#' @return A data frame with discounted costs, effectiveness and NMB.
#' @export
calculate_ce_out <- function(l_params_all, n_wtp = 100000, verbose= verbose){ # User defined
  with(as.list(l_params_all), {
    ########################### Process model inputs ###########################
    ## Number of cycles
    n_cycles   <- (n_age_max - n_age_init)/cycle_length # time horizon, number of cycles
    ## Age-specific transition probabilities to the Dead state
    # compute mortality rates
    v_r_S1Dage <- v_r_HDage * hr_S1 # Age-specific mortality rate in the Sick state 
    v_r_S2Dage <- v_r_HDage * hr_S2 # Age-specific mortality rate in the Sicker state 
    # transform rates to probabilities adjusting by cycle length
    p_HS1      <- rate_to_prob(r = r_HS1,  t = cycle_length) # constant annual probability of becoming Sick when Healthy conditional on surviving 
    p_S1H      <- rate_to_prob(r = r_S1H,  t = cycle_length) # constant annual probability of becoming Healthy when Sick conditional on surviving
    p_S1S2     <- rate_to_prob(r = r_S1S2, t = cycle_length) # constant annual probability of becoming Sicker when Sick conditional on surviving
    v_p_HDage  <- rate_to_prob(v_r_HDage,  t = cycle_length) # Age-specific mortality risk in the Healthy state 
    v_p_S1Dage <- rate_to_prob(v_r_S1Dage, t = cycle_length) # Age-specific mortality risk in the Sick state
    v_p_S2Dage <- rate_to_prob(v_r_S2Dage, t = cycle_length) # Age-specific mortality risk in the Sicker state
    
    ## Annual transition probability of becoming Sicker when Sick for treatment AB
    # Apply hazard ratio to rate to obtain transition rate of becoming Sicker when 
    # Sick for treatment B
    r_S1S2_trtAB <- r_S1S2 * hr_S1S2_trtAB
    # Transform rate to probability to become Sicker when Sick under treatment AB 
    # adjusting by cycle length conditional on surviving
    p_S1S2_trtAB <- rate_to_prob(r = r_S1S2_trtAB, t = cycle_length)
    
    ###################### Construct state-transition models ###################
    ## Initial state vector
    # All starting healthy
    v_m_init <- c(H = 1, S1 = 0, S2 = 0, D = 0) # initial state vector
    # Number of health states 
    n_states    <- length(v_m_init)
    # Health state names
    v_names_states <- names(v_m_init)
    
    ## Initialize cohort trace for age-dependent (ad) cSTM for strategies SoC and AB
    m_M_SoC <- matrix(0, 
                      nrow     = (n_cycles + 1), ncol = n_states, 
                      dimnames = list(0:n_cycles, v_names_states))
    # Store the initial state vector in the first row of the cohort trace
    m_M_SoC[1, ] <- v_m_init
    ## Initialize cohort trace for strategy AB
    # Structure and initial states are the same as for SoC
    m_M_strAB <- m_M_SoC 
    
    ## Initialize transition dynamics arrays which will capture transitions 
    ## from each state to another over time for strategy SoC
    a_A_SoC <- array(0,
                     dim      = c(n_states, n_states, n_cycles + 1),
                     dimnames = list(v_names_states, v_names_states, 0:n_cycles))
    # Set first slice of a_A_SoC with the initial state vector in its diagonal
    diag(a_A_SoC[, , 1]) <- v_m_init
    # Initialize transition-dynamics array for strategy AB
    # Structure and initial states are the same as for SoC
    a_A_strAB <- a_A_SoC
    
    ## Create transition arrays
    # Initialize 3-D array
    a_P_SoC <- array(0, dim   = c(n_states, n_states, n_cycles),
                     dimnames = list(v_names_states, v_names_states, 0:(n_cycles - 1)))
    ### Fill in array
    ## From H
    a_P_SoC["H", "H", ]   <- (1 - v_p_HDage) * (1 - p_HS1)
    a_P_SoC["H", "S1", ]  <- (1 - v_p_HDage) *      p_HS1
    a_P_SoC["H", "D", ]   <-      v_p_HDage
    ## From S1
    a_P_SoC["S1", "H", ]  <- (1 - v_p_S1Dage) * p_S1H
    a_P_SoC["S1", "S1", ] <- (1 - v_p_S1Dage) * (1 - (p_S1H + p_S1S2))
    a_P_SoC["S1", "S2", ] <- (1 - v_p_S1Dage) *               p_S1S2
    a_P_SoC["S1", "D", ]  <-      v_p_S1Dage
    ## From S2
    a_P_SoC["S2", "S2", ] <- 1 - v_p_S2Dage
    a_P_SoC["S2", "D", ]  <-     v_p_S2Dage
    ## From D
    a_P_SoC["D", "D", ]   <- 1
    
    ## Initialize transition probability matrix for strategy AB as a copy of SoC's
    a_P_strAB <- a_P_SoC
    
    ## Only need to update the probabilities involving the transition from Sick to Sicker, p_S1S2
    # From S1
    a_P_strAB["S1", "S1", ] <- (1 - v_p_S1Dage) * (1 - (p_S1H + p_S1S2_trtAB))
    a_P_strAB["S1", "S2", ] <- (1 - v_p_S1Dage) *               p_S1S2_trtAB
    
    ### Check if transition probability matrices are valid
    ## Check that transition probabilities are [0, 1]
    check_transition_probability(a_P_SoC, verbose = verbose)
    check_transition_probability(a_P_strAB, verbose = verbose)
    ### Check that all rows for each slice of the array sum to 1
    check_sum_of_transition_array(a_P_SoC, n_states = n_states, n_cycles = n_cycles, verbose = verbose)
    check_sum_of_transition_array(a_P_strAB, n_states = n_states, n_cycles = n_cycles, verbose = verbose)
    
    #### Run Markov model ####
    ## Iterative solution of age-dependent cSTM
    for (t in 1:n_cycles) {
      ## Fill in cohort trace
      # For SoC
      m_M_SoC[t + 1, ]   <- m_M_SoC[t, ]   %*% a_P_SoC[, , t]
      # For strategy AB 
      m_M_strAB[t + 1, ] <- m_M_strAB[t, ] %*% a_P_strAB[, , t]
      
      ## Fill in transition-dynamics array
      # For SoC
      a_A_SoC[, , t + 1]   <- diag(m_M_SoC[t, ])   %*% a_P_SoC[, , t]
      # For strategy AB
      a_A_strAB[, , t + 1] <- diag(m_M_strAB[t, ]) %*% a_P_strAB[, , t]
    }
    
    ## Store the cohort traces in a list
    l_m_M <- list(m_M_SoC,
                  m_M_strAB)
    names(l_m_M) <- v_names_str
    
    ## Store the transition array for each strategy in a list
    l_a_A <- list(a_A_SoC,
                  a_A_strAB)
    names(l_m_M) <- v_names_str
  
    
    #### State Rewards ####
    ## Vector of state utilities under strategy SoC
    # Your turn
    
    ## Vector of state costs under strategy SoC
    # Your turn
    
    ## Vector of state utilities under strategy AB
    # Your turn
    
    ## Vector of state costs under strategy AB
    # Your turn
    
    ## Store the vectors of state utilities for each strategy in a list 
    l_u   <- list(SoC = v_u_SoC,
                  AB  = v_u_strAB)
    ## Store the vectors of state cost for each strategy in a list 
    l_c   <- list(SoC = v_c_SoC,
                  AB  = v_c_strAB)
    
    # assign strategy names to matching items in the lists
    names(l_u) <- names(l_c) <- names(l_a_A) <- v_names_str
    
    ## create empty vectors to store total utilities and costs 
    v_tot_qaly <- v_tot_cost <- vector(mode = "numeric", length = n_str)
    names(v_tot_qaly) <- names(v_tot_cost) <- v_names_str
    
    ## Number of cycles
    n_cycles <- (n_age_max - n_age_init)/cycle_length # time horizon, number of cycles
    
    ## Discount weight for costs and effects
    v_dwc  <- 1 / ((1 + d_c * cycle_length) ^ (0:n_cycles))
    v_dwe  <- 1 / ((1 + d_e * cycle_length) ^ (0:n_cycles))
    
    ## Within-cycle correction (WCC) using Simpson's 1/3 rule
    v_wcc <- darthtools::gen_wcc(n_cycles = n_cycles, 
                                 method = "Simpson1/3") # vector of wcc
    
    #### Loop through each strategy and calculate total utilities and costs ####
    for (i in 1:n_str) { # i <- 1
      v_u_str <- l_u[[i]]   # select the vector of state utilities for the i-th strategy
      v_c_str <- l_c[[i]]   # select the vector of state costs for the i-th strategy
      a_A_str <- l_a_A[[i]] # select the transition array for the i-th strategy
      
      #### Array of state rewards ####
      # Create transition matrices of state utilities and state costs for the i-th strategy 
      m_u_str   <- matrix(v_u_str, nrow = n_states, ncol = n_states, byrow = T)
      m_c_str   <- matrix(v_c_str, nrow = n_states, ncol = n_states, byrow = T)
      # Expand the transition matrix of state utilities across cycles to form a transition array of state utilities
      a_R_u_str <- array(m_u_str, 
                         dim      = c(n_states, n_states, n_cycles + 1),
                         dimnames = list(v_names_states, v_names_states, 0:n_cycles))
      # Expand the transition matrix of state costs across cycles to form a transition array of state costs
      a_R_c_str <- array(m_c_str, 
                         dim      = c(n_states, n_states, n_cycles + 1),
                         dimnames = list(v_names_states, v_names_states, 0:n_cycles))
      
      #### Apply transition rewards ####  
      # Apply disutility due to transition from H to S1
      a_R_u_str["H", "S1", ]      <- a_R_u_str["H", "S1", ]       - du_HS1
      # Add transition cost per cycle due to transition from H to S1
      a_R_c_str["H", "S1", ]      <- a_R_c_str["H", "S1", ]       + ic_HS1
      # Add transition cost  per cycle of dying from all non-dead states
      a_R_c_str[-n_states, "D", ] <- a_R_c_str[-n_states, "D", ] + ic_D
      
      #### Expected QALYs and Costs for all transitions per cycle ####
      # QALYs = life years x QoL
      # Note: all parameters are annual in our example. In case your own case example is different make sure you correctly apply .
      a_Y_c_str <- a_A_str * a_R_c_str
      a_Y_u_str <- a_A_str * a_R_u_str 
      
      #### Expected QALYs and Costs per cycle ####
      ## Vector of QALYs and Costs
      v_qaly_str <- apply(a_Y_u_str, 3, sum) # sum the proportion of the cohort across transitions 
      v_cost_str <- apply(a_Y_c_str, 3, sum) # sum the proportion of the cohort across transitions
      
      #### Discounted total expected QALYs and Costs per strategy and apply half-cycle correction if applicable ####
      ## QALYs
      v_tot_qaly[i] <- t(v_qaly_str) %*% (v_dwe * v_wcc)
      ## Costs
      v_tot_cost[i] <- t(v_cost_str) %*% (v_dwc * v_wcc)
    }
    
    ## Vector with discounted net monetary benefits (NMB)
    v_nmb <- v_tot_qaly * n_wtp - v_tot_cost
    
    ## data.frame with discounted costs, effectiveness and NMB
    df_ce <- data.frame(Strategy = v_names_str,
                        Cost     = v_tot_cost,
                        Effect   = v_tot_qaly,
                        NMB      = v_nmb)
    
    return(df_ce)
  }
  )
}

#------------------------------------------------------------------------------#
####             Generate a PSA input parameter dataset                     ####
#------------------------------------------------------------------------------#
#' Generate parameter sets for the probabilistic sensitivity analysis (PSA)
#'
#' \code{generate_psa_params} generates a PSA dataset of the parameters of the 
#' cost-effectiveness analysis.
#' @param n_sim Number of parameter sets for the PSA dataset
#' @param seed Seed for the random number generation
#' @return A data.frame with a PSA dataset of he parameters of the 
#' cost-effectiveness analysis
#' @export
generate_psa_params <- function(n_sim = 1000, seed = 071818){
  set.seed(seed) # set a seed to be able to reproduce the same results
  df_psa <- data.frame(
    # Transition probabilities (per cycle)
    r_HS1    = rgamma(n_sim, shape = 30, rate = 170 + 30),     # constant rate of becoming Sick when Healthy conditional on surviving
    r_S1H    = rgamma(n_sim, shape = 60, rate = 60 + 60),      # constant rate of becoming Healthy when Sick conditional on surviving
    r_S1S2   = rgamma(n_sim, shape = 84, rate = 716 + 84),     # constant rate of becoming Sicker when Sick conditional on surviving
    hr_S1    = rlnorm(n_sim, meanlog = log(3), sdlog = 0.01),  # hazard ratio of death in Sick vs healthy
    hr_S2    = rlnorm(n_sim, meanlog = log(10), sdlog = 0.02), # hazard ratio of death in Sicker vs healthy 
    
    # Effectiveness of treatment AaB 
    hr_S1S2_trtAB = rlnorm(n_sim, meanlog = log(0.6), sdlog = 0.02), # hazard ratio of becoming Sicker when Sick under treatment AB
    
    ## State rewards
    # Costs
    c_H     = rgamma(n_sim, shape = 100,   scale = 20),   # cost of remaining one cycle in state H
    c_S1    = rgamma(n_sim, shape = 177.8, scale = 22.5), # cost of remaining one cycle in state S1
    c_S2    = rgamma(n_sim, shape = 225,   scale = 66.7), # cost of remaining one cycle in state S2
    c_trtAB = rgamma(n_sim, shape = 156.3, scale = 160),  # cost of treatment AB (per cycle)
    c_D     = 0,                                          # cost of being in the death state
    
    # Utilities
    u_H     = rbeta(n_sim, shape1 = 200, shape2 = 3),     # utility when healthy
    u_S1    = rbeta(n_sim, shape1 = 130, shape2 = 45),    # utility when sick
    u_S2    = rbeta(n_sim, shape1 = 230, shape2 = 230),   # utility when sicker
    u_D     = 0,                                          # utility when dead
    u_trtAB = rbeta(n_sim, shape1 = 300, shape2 = 15),    # utility when being treated with AB
    
    ## Transition rewards
    du_HS1  = rbeta(n_sim, shape1 = 11,  shape2 = 1088),  # disutility when transitioning from Healthy to Sick
    ic_HS1  = rgamma(n_sim, shape = 25,  scale = 40),     # increase in cost when transitioning from Healthy to Sick
    ic_D    = rgamma(n_sim, shape = 100, scale = 20)      # increase in cost when dying
  )
  return(df_psa)
}
