# Graduate seminar on operations research
# Homework 20

# Add CABG as a third treatment action. 
# The treatment is much more efficient and expensive compared to PTCA. 
# How does the optimal policy change with your chosen values?



#install.packages("pomdp")
library("pomdp")
# documentation at http://search.r-project.org/R/R/library/pomdp/html/solve_POMDP.html

## This model assumes the patient visits a physician once a month.
## One action can be done each time.

ihd_cost <- -30 # cost of one month with severe IHD

# costs of the actions
PTCA_cost <- -50
medication_cost <- -5
angiogram_cost <- -5
stresstest_cost <- -1

# Transition probabilities each month without treatment
waiting_transitions <- matrix(c(0.9, 0.1, 0.01, 0.99), nrow = 2 , byrow = TRUE)



# define the problem
IHD_diagnostics <- POMDP(

  states = c("no_IHD", "IHD"),
  actions = c("stress-test", "angiogram", "medication", "PTCA", "wait"),
  observations = c("no_IHD", "IHD"),
  discount = 0.95, # assuming the intervals are always the same
  
  transition_prob = list(
    "stress-test" = waiting_transitions , 
    "angiogram" = waiting_transitions,
    "wait" = waiting_transitions,
    "medication" = matrix(c( 
      0.9, 0.1,
      0.3,   0.7), nrow = 2 , byrow = TRUE) , 
    "PTCA" = matrix(c( 
      0.99, 0.01,
      0.9,   0.1), nrow = 2 , byrow = TRUE) 
  ),
  
  observation_prob = list(
    "stress-test" = matrix(c( 
      0.85, 0.15,
      0.2,   0.8), nrow = 2 , byrow = TRUE), 
    "angiogram" = matrix(c( 
      1, 0,
      0,   1), nrow = 2 , byrow = TRUE),
    "medication" = "uniform",
    "PTCA" = "uniform",
    "wait" = "uniform"
  ) ,
  
  reward = rbind(
    R_("medication",  "*",  "IHD",      "*", medication_cost + ihd_cost),
    R_("medication",  "*",  "no_IHD",   "*", medication_cost),
    R_("PTCA",        "*",  "IHD",      "*", PTCA_cost + ihd_cost),
    R_("PTCA",        "*",  "no_IHD",   "*", PTCA_cost),
    R_("angiogram",   "*",  "IHD",      "*", angiogram_cost + ihd_cost),
    R_("angiogram",   "*",  "no_IHD",   "*", angiogram_cost),
    R_("stress-test", "*",  "IHD",      "*", stresstest_cost + ihd_cost),
    R_("stress-test", "*",  "no_IHD",   "*", stresstest_cost),
    R_("wait",        "*",  "IHD",      "*", ihd_cost),
    R_("wait",        "*",  "no_IHD",   "*", 0)
  ) ,
  
  start = "uniform",
)



# solve the problem
sol <- solve_POMDP(IHD_diagnostics)
sol$solution
sol$solution$alpha

# visualize the solution
plot_policy_graph(sol)
plot_value_function(sol)