# 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 exponential distribution with same mean survival as Weibull
# distribution for Diseased to Dead
mean_surv <- lambda_surv ^ (-1 / gamma_surv) * gamma(1 + 1 / 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), 3), list(cycle = cycles, transition = c("Healthy->Diseased", "Healthy->Dead", "Diseased->Dead")))
transition_probs[, 1:2] <- healthy_haz / rowSums(healthy_haz) * (1 - exp(-rowSums(healthy_haz)))
transition_probs[, 3] <- 1 - exp(-rate_surv)

# Run cohort simulation
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]
}