# We will use cycles from 0 to 100
cycles <- 0:100

# Create an array for the state membership
state_membership <- array(NA_real_, c(length(cycles), 3), list(cycle = cycles, state = c("Healthy", "Diseased", "Dead")))
state_membership[1, ] <- c(1, 0, 0)

# Define the survival model parameters
# - Healthy to Diseased
lambda_inc  <- 2e-5
gamma_inc   <- 2.5
# - Healthy to Dead
a_general   <- 0.2
b_general   <- 1e-7
# - Diseased to Dead
lambda_surv <- 0.2
gamma_surv  <- 0.5

# Calculate the hazard rates for the transitions out of the Healthy state
healthy_haz <- array(NA_real_, c(length(cycles), 2), list(cycle = cycles, to = c("Diseased", "Dead")))
healthy_haz[, 1] <- lambda_inc * ((cycles + 1) ^ gamma_inc - cycles ^ gamma_inc)
healthy_haz[, 2] <- b_general / a_general * (exp(a_general * (cycles + 1)) - exp(a_general * cycles))

# Create an array for transition probabilities in each cycle
transition_probs <- array(NA_real_, c(length(cycles), 3), list(cycle = cycles, transition = c("Healthy->Diseased", "Healthy->Dead", "Diseased->Dead")))

# The transition probabilities for the Healthy state can be calculated already
transition_probs[, 1:2] <- healthy_haz / rowSums(healthy_haz) * (1 - exp(-rowSums(healthy_haz)))

# Create an array for the moments of the overall CSSTD in each cycle
csstd_moments <- array(NA_real_, c(length(cycles), 5), list(cycle = cycles, moment = c("1st raw (mean)", "2nd raw", "3rd raw", "2nd central (variance)", "3rd central")))

# Create an array for the weights of the two components in the CSSTD
component_weights <- array(NA_real_, c(length(cycles), 2), list(cycle = cycles, component = c("Remaining in state", "Entering state")))

# Create an array for the moments for the first component (remaining in the state)
remaining_moments <- array(NA_real_, c(length(cycles), 6), list(cycle = cycles, moment = c("1st raw (mean)", "1st around 1", "2nd around 1", "3rd around 1", "2nd central (variance)", "3rd central")))

# Create a vector for the probability for remaining in the Diseased state
ephi <- rep(NA_real_, length(cycles))

# Create an array for rho
rho <- array(NA_real_, c(length(cycles), 3), list(cycle = cycles, derivative = 0:2))

# Create an expression for phi(s), the conditional probability of staying in the state one more cycle
# Note: We aren't evaluating this yet - it is just symbolic
phi <- expression( exp(-lambda_surv * ((s + 1) ^ gamma_surv - s ^ gamma_surv)) )

# Create expressions for the first two derivatives of phi(s)
dphi  <- D(  phi, "s")
d2phi <- D( dphi, "s")
d3phi <- D(d2phi, "s")

# Initialise the model
csstd_moments[1, ] <- 0
ephi[1] <- eval(phi, list(s = 0))
transition_probs[1, 3] <- 1 - ephi[1]
rho[1, 1] <- 1
rho[1, 2:3] <- 0

# Iterate through the cycles
for (i in 2:length(cycles)) {
  
  # Update state membership
  state_membership[i, 1] <- state_membership[i - 1, 1] * (1 - sum(transition_probs[i - 1, 1:2]))
  state_membership[i, 2] <- state_membership[i - 1, 1] * transition_probs[i - 1, 1] +
    state_membership[i - 1, 2] * (1 - transition_probs[i - 1, 3])
  state_membership[i, 3] <- state_membership[i - 1, 1] * transition_probs[i - 1, 2] +
    state_membership[i - 1, 2] * transition_probs[i - 1, 3] +
    state_membership[i - 1, 3]
  
  # Calculate weights
  component_weights[i, 2] <- state_membership[i - 1, 1] * transition_probs[i - 1, 1] / state_membership[i, 2]
  component_weights[i, 1] <- 1 - component_weights[i, 2]
  
  # Calculate moments for those remaining in the state (since the previous cycle)
  remaining_moments[i, 1] <- 1 + csstd_moments[i - 1, "1st raw (mean)"] * rho[i - 1, 1] +
    (csstd_moments[i - 1, "2nd raw"] - csstd_moments[i - 1, "1st raw (mean)"] ^ 2) * rho[i - 1, 2] +
    (csstd_moments[i - 1, "3rd raw"] - 2 * csstd_moments[i - 1, "1st raw (mean)"] * csstd_moments[i - 1, "2nd raw"] + csstd_moments[i - 1, "1st raw (mean)"] ^ 3) * rho[i - 1, 3] / 2
  remaining_moments[i, 2] <- remaining_moments[i, 1] - 1
  remaining_moments[i, 3] <- csstd_moments[i - 1, "2nd raw"] * rho[i - 1, 1] +
    (csstd_moments[i - 1, "3rd raw"] - csstd_moments[i - 1, "1st raw (mean)"] * csstd_moments[i - 1, "2nd raw"]) * rho[i - 1, 2]
  remaining_moments[i, 4] <- csstd_moments[i - 1, "3rd raw"] * rho[i - 1, 1]
  remaining_moments[i, 5] <- remaining_moments[i, 3] - (1 - remaining_moments[i, 1]) ^ 2
  remaining_moments[i, 6] <- remaining_moments[i, 4] + 3 * remaining_moments[i, 3] * (1 - remaining_moments[i, 1]) - 2 * (1 - remaining_moments[i, 1]) ^ 3
  
  # Calculate overall CSSTD moments
  csstd_moments[i, 1] <- component_weights[i, 1] * remaining_moments[i, 1] + component_weights[i, 2] * 0.5
  temp <- remaining_moments[i, 1] - csstd_moments[i, 1]
  csstd_moments[i, 4] <- component_weights[i, 1] * (temp^2 + remaining_moments[i, 5]) + component_weights[i, 2] * (0.5 - csstd_moments[i, 1]) ^ 2
  csstd_moments[i, 5] <- component_weights[i, 1] * (temp^3 + 3 * remaining_moments[i, 5] * temp + remaining_moments[i, 6]) + component_weights[i, 2] * (0.5 - csstd_moments[i, 1])^3
  csstd_moments[i, 2] <- csstd_moments[i, 1] ^ 2 + csstd_moments[i, 4]
  csstd_moments[i, 3] <- csstd_moments[i, 1] ^ 3 + 3 * csstd_moments[i, 4] * csstd_moments[i, 1] + csstd_moments[i, 5]
  
  # Calculate E[phi(s)]
  mu <- csstd_moments[i, 1]
  phi_eval   <- eval(  phi, list(s = mu))
  dphi_eval  <- eval( dphi, list(s = mu))
  d2phi_eval <- eval(d2phi, list(s = mu))
  d3phi_eval <- eval(d3phi, list(s = mu))
  
  ephi[i] <- phi_eval + d2phi_eval / 2 * csstd_moments[i, "2nd central (variance)"] + d3phi_eval / 6 * csstd_moments[i, "3rd central"]
  transition_probs[i, 3] <- 1 - ephi[i]
  
  # Calculate rho and derivatives
  rho[i, 1] <- phi_eval / ephi[i]
  rho[i, 2] <- dphi_eval / ephi[i]
  rho[i, 3] <- d2phi_eval / ephi[i]
}