library(dplyr)
library(diagram)
library(ggplot2)
library(xlsx)
library(backports)
#devtools::install_github("pierucci/heemod", force = TRUE)
library(heemod)

param <- define_parameters(

  #Transition probabilities:
  
  #Arm1
  A1toR = 0.095592,
  A1toA = 0.040113,
  A1toD = 0.157010,
  A1toS = 1-A1toR-A1toA-A1toD,
  A1toC = 1-A1toA-A1toD,
  
  #Arm2:
  A2toR = 0.183585,
  A2toA = 0.097994,
  A2toD = 0.202784,
  A2toS = 1-A2toR-A2toA-A2toD,
  A2toC = 1-A2toA-A2toD,
  
  
  #After A:
  AAlive = 0.805927,
  AD = 1-AAlive,
  
  
  #Utilities for Quality of Life:
  u_arm1 = .56,
  u_arm2 = .40,
  u_S1 = .95,
  u_R = .57,
  u_S2 = .90,
  u_aalive = .54,
  u_d = 0
)

#Define Transition Matrix:

#Arm1:

trans_arm1 <- define_transition(
  state_names = c("I", "S1", "R", "S2", "C", "A", "D"),
  0, A1toS, A1toR, 0, 0, A1toA, A1toD,
  0, A1toS, A1toR, 0, 0, A1toA, A1toD,
  0, 0, 0, A1toS, A1toR, A1toA, A1toD,
  0, 0, 0, A1toS, A1toR, A1toA, A1toD,
  0, 0, 0, 0, A1toC, A1toA, A1toD,
  0, 0, 0, 0, 0, AAlive, AD,
  0, 0, 0, 0, 0, 0, 1
)

#Arm2:

trans_arm2 <- define_transition(
  state_names = c("I", "S1", "R", "S2", "C", "A", "D"),
  0, A2toS, A2toR, 0, 0, A2toA, A2toD,
  0, A2toS, A2toR, 0, 0, A2toA, A2toD,
  0, 0, 0, A1toS, A1toR, A1toA, A1toD,
  0, 0, 0, A1toS, A1toR, A1toA, A1toD,
  0, 0, 0, 0, A1toC, A1toA, A1toD,
  0, 0, 0, 0, 0, AAlive, AD,
  0, 0, 0, 0, 0, 0, 1
)


#Define Stages:

I <- define_state(
  total.cost=100 + 500,
  utility=(u_arm1)
)

S1 <- define_state(
  total.cost=100,
  utility=(u_S1)
)

R <- define_state(
  total.cost=100 + 500,
  utility=(u_R)
)

S2 <- define_state(
  total.cost=100,
  utility=(u_S2)
)

C <- define_state(
  total.cost=100,
  utility=(u_arm2)
)

A <- define_state(
  total.cost=100,
  starting_values=define_starting_values(total.cost=1000),
  utility=(u_aalive)
)

D <- define_state(
  total.cost=0,
  utility=u_d
)

Arm2 <- define_state(
  total.cost=100,
  utility=(u_arm2)
)


#Define strategy: 


strat.one <- define_strategy(
  transition = trans_arm1,
  I = I, S1 = S1, R = R, S2 = S2, C = C, A = A, D = D
)

strat.two <- define_strategy(
  transition = trans_arm2,
  I = Arm2, S1 = Arm2, R = R, S2 = S2, C = C, A = A, D = D
)


#Run the model

mod.basic <- run_model(
  treatment = strat.one,
  control = strat.two,
  parameters = param,
  cycles = 6,
  cost = total.cost,
  effect = utility,
  init = c(100, 0, 0, 0, 0, 0, 0),
  method="end"
)

summary(mod.basic)

counts <- get_counts(mod.basic)

values <- get_values(mod.basic)

values2 <- filter(values,.strategy_names == "treatment")