# 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), 5), list(cycle = cycles, state = c("Healthy", "DiseasedA", "DiseasedB", "DiseasedC", "Dead")))
state_membership[1, ] <- c(1, 0, 0, 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))

# Fit a phase-type distribution to the Weibull distribution for disease survival
fit <- mapfit::phfit.density(
  ph = mapfit::cf1(3),
  f = function(x) lambda_surv * gamma_surv * x ^ (gamma_surv - 1) * exp(-lambda_surv * x ^ gamma_surv))
Q <- Matrix::as.matrix(fit$model@Q)
P <- expm::expm(Q)

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

# The transition probabilities for the Healthy state can be calculated already
healthy_to_diseased <- healthy_haz[, 1] / rowSums(healthy_haz) * (1 - exp(-rowSums(healthy_haz)))
transition_probs[, 1:3] <- sapply(1:3, function(i) healthy_to_diseased * fit$model@alpha[i])
transition_probs[, 4] <- healthy_haz[, 2] / rowSums(healthy_haz) * (1 - exp(-rowSums(healthy_haz)))
transition_probs[, 5:10] <- t(sapply(cycles, function(i) P[upper.tri(P, diag = TRUE)]))

# Run the cohort simulation
for (i in 2:length(cycles)) {
  TP <- matrix(0, 5, 5)
  TP[1, 5] <- transition_probs[i - 1, 4]
  TP[1, 2:4] <- transition_probs[i - 1, 1:3]
  TP[1, 1] <- 1 - sum(TP[1, 2:5])
  TP[(row(TP) %in% (2:4)) & (col(TP) %in% (2:4) & (row(TP) <= col(TP)))] <- transition_probs[i - 1, 5:10]
  TP[2:5, 5] <- 1 - rowSums(TP[2:5, 1:4])
  
  state_membership[i, ] <- t(TP) %*% state_membership[i - 1, ]
}

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