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

# We use 10 states for Diseased (9 tunnel and 1 with self-loop)
n_tunnel <- 98
states <- c("Healthy", paste0("Diseased", sprintf("%02i", 1:(n_tunnel+1))), "Dead")

# Create an array for the state membership
state_membership <- array(NA_real_, c(length(cycles), length(states)), list(cycle = cycles, state = states))
state_membership[1, ] <- c(1, rep(0, length(states) - 1))

# 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 exponential distribution with same mean survival as conditional
# Weibull distribution for Diseased to Dead given survived to 10th state
upper_incomplete_gamma <- function(a, x) pgamma(x, a, lower = FALSE) * gamma(a)
mean_surv <- exp(lambda_surv * n_tunnel ^ gamma_surv) * lambda_surv ^ (-1 / gamma_surv) * upper_incomplete_gamma(1 + 1 / gamma_surv, lambda_surv * n_tunnel ^ gamma_surv)
rate_surv <- 1 / mean_surv

# 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), 2 + n_tunnel + 1),
  list(
    cycle = cycles,
    transition = c(
      "Healthy->Diseased01",
      "Healthy->Dead",
      paste0("Diseased", sprintf("%02i", 1:(n_tunnel)), "->Diseased", sprintf("%02i", 2:(n_tunnel + 1))),
      paste0("Diseased", sprintf("%02i", n_tunnel + 1), "->Diseased", sprintf("%02i", n_tunnel + 1)))))
transition_probs[, 1:2] <- healthy_haz / rowSums(healthy_haz) * (1 - exp(-rowSums(healthy_haz)))
for (t in 1:n_tunnel) {
  transition_probs[, t + 2] <- exp(-lambda_surv * (t ^ gamma_surv - (t - 1) ^ gamma_surv))
}
transition_probs[, n_tunnel + 3] <- exp(-rate_surv)

# Run cohort simulation
for (i in 2:length(cycles)) {
  # Update state membership
  state_membership[i, "Healthy"] <- state_membership[i - 1, "Healthy"] * (1 - transition_probs[i - 1, "Healthy->Diseased01"] - transition_probs[i - 1, "Healthy->Dead"])
  state_membership[i, "Diseased01"] <- state_membership[i - 1, "Healthy"] * transition_probs[i - 1, "Healthy->Diseased01"]
  for (t in 1:n_tunnel) {
    from <- paste0("Diseased", sprintf("%02i", t))
    to   <- paste0("Diseased", sprintf("%02i", t + 1))
    state_membership[i, to] <- state_membership[i - 1, from] * transition_probs[i - 1, paste0(from, "->", to)]
  }
  last <- paste0("Diseased", sprintf("%02i", n_tunnel + 1))
  state_membership[i, last] <- state_membership[i - 1, last] * transition_probs[i - 1, paste0(last, "->", last)] + state_membership[i, last]
  state_membership[i, "Dead"] <- 1 - sum(state_membership[i, 1:(n_tunnel+2)])
}

summarised_state_membership <- cbind(state_membership[, "Healthy"], rowSums(state_membership[, 2:(n_tunnel + 2)]), state_membership[, "Dead"])
dimnames(summarised_state_membership) <- list(cycle = cycles, state = c("Healthy", "Diseased", "Dead"))