function for value of information calculation

[zelihay]

Hello:)

Thank you for the pyAgrum package, which I am using extensively for my dissertation.

I am currently working on creating a function to calculate the value of information (VoI), which supposed to be straightforward. I created two versions of the functions; 1) adding or removing the arc between the chance node and the decision node, and 2) calculating marginal probabilities and expected utility values with evidence, and I applied them on a toy model. They calculated voi correctly and gave the same result, which is the same with manual calculation.

However, when I applied them to a relatively large model, the VoI often turns out to be zero or negative. I am in doubt whether those chance nodes are insignificant and the VoI calculations are correct, or if there is an issue with the functions.

Is there anyone who might be interested take a look at the codes, below?

Many thanks:)
Zeliha

import pyAgrum as gum
import pyAgrum.lib.notebook as gnb 
import random
import math

id_model=gum.InfluenceDiagram()
Take_Umbrella=id_model.addDecisionNode(gum.LabelizedVariable("Take_Umbrella", "Take_Umbrella", ["Yes","No"])) #the decision node
# chance nodes
Weather=id_model.addChanceNode(gum.LabelizedVariable("Weather", "Weather", ["rain","no_Rain"]))
Forecast=id_model.addChanceNode(gum.LabelizedVariable("Forecast", "Forecast", ["Sunny","Cloudy","rainy"]))

Happiness=id_model.addUtilityNode(gum.LabelizedVariable("Happiness", "Happiness", 1))#the utility node

id_model.addArc(Take_Umbrella,Happiness)
id_model.addArc(Weather,Forecast)
id_model.addArc(Weather,Happiness)

gnb.showInfluenceDiagram(id_model)

#CPTs

id_model.cpt(Weather).fillWith([0.3, 0.7])
id_model.cpt(Forecast)[{"Weather":"rain"}]=[0.05,0.15,0.80]
id_model.cpt(Forecast)[{"Weather":"no_Rain"}]=[0.5,0.3,0.2]

# utilities

id_model.utility(Happiness)[{'Weather': "rain", 'Take_Umbrella':"Yes" }]=80
id_model.utility(Happiness)[{'Weather': "rain", 'Take_Umbrella':"No" }]=0
id_model.utility(Happiness)[{'Weather': "no_Rain", 'Take_Umbrella':"Yes" }]=90
id_model.utility(Happiness)[{'Weather': "no_Rain", 'Take_Umbrella':"No" }]=100


### METHOD #1 ADDING ARC TO THE MODEL, AND CALCULATING VOI AFTERWARDS

##function for baseline eu calculation
def calculate_baseline_eu(model):
        limid_solver = gum.ShaferShenoyLIMIDInference(model)
        limid_solver.makeInference()
        baseline_meu = limid_solver.MEU()
        return baseline_meu.get("mean", float("-inf")) if isinstance(baseline_meu, dict) else baseline_meu

##function for voi calculation 
def calculate_voi_by_checking_arc_first(model, chance_node_name):
    decision_node_id = next(node for node in model.nodes() if model.isDecisionNode(node))
    decision_node_name = model.variable(decision_node_id).name()
    
    chance_node_id = model.idFromName(chance_node_name)
    
    arc_exists = model.existsArc(chance_node_id, decision_node_id)#check if an arc exist between the given chance node and the decision node
    
    if arc_exists:#if the arc exist remove it to calculate baseline eu
        print(f"\nArc from {chance_node_name} to {decision_node_name} exists. Removing it temporarily to calculate baseline EU.")
        model.eraseArc(chance_node_id, decision_node_id)
        baseline_eu = calculate_baseline_eu(model)
        print(f"\nBaseline Expected Utility (without arc): {baseline_eu}")
        
        print(f"\nRestoring the previously removed arc from {chance_node_name} to {decision_node_name}")
        model.addArc(chance_node_id, decision_node_id)
    else:#if the arc doesn't exist calculate the baseline eu first, then add the arc to calculate eu with evidence
        print(f"\nNo existing arc from {chance_node_name} to {decision_node_name}. Calculating baseline EU before adding the arc.")
        baseline_eu = calculate_baseline_eu(model)
        print(f"\nBaseline Expected Utility (without arc): {baseline_eu}")
        # Add the arc
        print(f"\nAdding arc from {chance_node_name} to {decision_node_name}")
        model.addArc(chance_node_id, decision_node_id)
    
    updated_eu = calculate_baseline_eu(model)#calculate eu with the arc
    print(f"\nUpdated Expected Utility (with arc from {chance_node_name} to {decision_node_name}): {updated_eu}")
    model.eraseArc(chance_node_id, decision_node_id)
    # Calculate voi
    eu_difference = updated_eu - baseline_eu
    print(f"\nDifference in Expected Utility (VoI for adding arc from {chance_node_name} to {decision_node_name}): {eu_difference}")

    return baseline_eu, updated_eu, eu_difference

## call the function for Forecast:
#calculate_voi_by_checking_arc_first(id_model, "Forecast")#---> PLEASE COMMENT OUT THIS BEFORE CALLING THE METHOD #2

### METHOD #2 CALCULATING VOI WITH MARGINAL PROBABILITIES AND EXPECTED UTILITIES WITH THE EVIDENCE OF THE CHANCE NODE(FORECAST)

def calculate_voi_with_marginal_probabilities(model, chance_node_name):
    #Calculating baseline EU
    baseline_eu = calculate_baseline_eu(model)
    
    #Determining if the given chance node has parents
    chance_node_id = model.idFromName(chance_node_name)
    has_parents = len(model.parents(chance_node_id)) > 0
    chance_node_states = model.variable(chance_node_id).labels()

    #If the chance node has parents get marginal probabilities by posterior function
    marginal_probs = {}
    if has_parents:
        # Get marginal probabilities by inference
        print(f"\nCalculating marginal probabilities for '{chance_node_name}' with parents.")
        limid = gum.ShaferShenoyLIMIDInference(model)
        limid.makeInference()
        marginal_distribution = limid.posterior(chance_node_id)
        for i, prob in enumerate(marginal_distribution.tolist()):
            state_label = chance_node_states[i]
            marginal_probs[state_label] = prob
            print(f"{state_label}: Probability = {prob}")
    else:
        # If the chance node does not have parents use prior probabilities directly
        print(f"\nUsing prior probabilities from CPT for '{chance_node_name}' (no parents).")
        cpt = model.cpt(chance_node_id)
        for i, prob in enumerate(cpt.tolist()):
            state_label = chance_node_states[i]
            marginal_probs[state_label] = prob
            print(f"{state_label}: Probability = {prob}")
    
    # Calculate expected utility for each state of given chance node
    expected_utilities = {}
    limid = gum.ShaferShenoyLIMIDInference(model)
    for state in chance_node_states:
        limid.setEvidence({chance_node_name: state})
        limid.makeInference()
        expected_utility = limid.MEU()
        expected_utilities[state] = float(expected_utility.get("mean", float('-inf')) if isinstance(expected_utility, dict) else expected_utility)
        limid.setEvidence({})  # Clear evidence for the next iteration
    
    # Calculate expected utility with perfect information
    weighted_eu_with_perfect_info = sum(
        marginal_probs[state] * expected_utilities[state] for state in chance_node_states
    )
    
    # Calculate the VoI
    voi = weighted_eu_with_perfect_info - baseline_eu    
   
    print(f"\nBaseline Expected Utility (without evidence): {baseline_eu}")
    print(f"Weighted Expected Utility (with perfect information on '{chance_node_name}'): {weighted_eu_with_perfect_info}")
    print(f"VoI for '{chance_node_name}': {voi}")    
    return voi
    
#calculate_voi_with_marginal_probabilities(id_model, "Forecast")

[phwuil]

Hi Zeliha (@zelihay),

Can you show an example where VOI seems to be incorrect?
For now, I think that the second method (calculate_voi_with_marginal_probabilities) should be more general

Let me show you another way to implement it (more compact) :

def VOI(model,node_name,verbose=True):
  def getMEU(evidence):
    limid = gum.ShaferShenoyLIMIDInference(model)
    limid.setEvidence(evidence)
    limid.makeInference()
    meu=limid.MEU()['mean']
    if verbose:
      print(f"Inference with {evidence=} : {meu}")
    return meu
    
  limid = gum.ShaferShenoyLIMIDInference(model)
  limid.makeInference()
  posterior=limid.posterior(node_name)
  baseline=limid.MEU()['mean']
  if verbose:
    print(f"{baseline=}")

  
  averageMEU=posterior.expectedValue(getMEU)
  if verbose:
    print(f"{averageMEU=}")
  voi=averageMEU-baseline
  if verbose:
    print(f"{voi=}")
  
  return voi

[zelihay]

Hi Mr. Wuillemin,

Sorry for the late reply.

Thanks a lot for your time to provide the function. It is much shorter and more to the point. :) In case you would like to see the model that gives 0 voi for all chance nodes, here are the codes:
(Note: After writing the previous message, I played around with the parameters of the utility functions, and it turns out that higher increments' values give positive VOI values.)

Thanks again.

Zeliha

import pyAgrum as gum
import numpy as np
import itertools  
#------------------------------------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------DOCTOR MODEL------------------------------------------------------------------------------------------
#------------------------------------------------------------------------------------------------------------------------------------------------------

import pyAgrum as gum

doctor_model = gum.InfluenceDiagram()
    
age = doctor_model.add(gum.LabelizedVariable("Age", "Age of the patient", 3))  

smoking = doctor_model.add(gum.LabelizedVariable("Smoking", "Smoking", 2))  # No, Yes
goiter_size = doctor_model.add(gum.LabelizedVariable("GoiterSize", "Goiter Size", 2))  # Small, Large


Eye_Disease_Status = doctor_model.add(gum.LabelizedVariable("Eye_Disease_Status", "Eye_Disease_Status", 3))  # low, Moderate, severe
tyroid_eye_disease = doctor_model.add(gum.LabelizedVariable("Tyroid_Eye_Disease", "Tyroid_Eye_Disease", 2))  # No, Yes

Thyroid_Function_Status = doctor_model.add(gum.LabelizedVariable("Thyroid_Function_Status", "Thyroid_Function_Status", 3)) #Normal , Overactive, Underactive 

Duration_of_Treatment = doctor_model.add(gum.LabelizedVariable("Duration_of_Treatment", "Duration_of_Treatment", 3)) #Short , Medium, Long 

cardiovascular_risk = doctor_model.add(gum.LabelizedVariable("CardioRisk", "Cardiovascular Risk", 2))  # No, Yes
General_Health_Status = doctor_model.add(gum.LabelizedVariable("General_Health_Status", "General Health Status", 3))  # Good, Average, Poor

TSH_Level = doctor_model.add(gum.LabelizedVariable("TSH_Level", "General TSH_Level Status", 3)) 

risk_of_relapse = doctor_model.add(gum.LabelizedVariable("RelapseRisk", "Risk of Relapse", 2))  # Low, High

severity_hypothyroidism = doctor_model.add(gum.LabelizedVariable("Severity_Hypothyroidism", "Severity_Hypothyroidism", 2))  # No, Yes

treatment = doctor_model.addDecisionNode(gum.LabelizedVariable("Treatment", "Chosen Treatment", ["ATD", "RAI", "Surgery"]))  # ATD, RAI, Surgery

satisfaction = doctor_model.addUtilityNode(gum.LabelizedVariable("Satisfaction", "Patient Satisfaction", 1))  # Low, Medium, High

doctor_model.addArc(age, General_Health_Status)
doctor_model.addArc(age, goiter_size)

doctor_model.addArc(General_Health_Status,cardiovascular_risk)
doctor_model.addArc(smoking,General_Health_Status)
doctor_model.addArc(smoking,cardiovascular_risk)
doctor_model.addArc(age,cardiovascular_risk)

doctor_model.addArc(goiter_size, risk_of_relapse)
doctor_model.addArc(goiter_size, severity_hypothyroidism)
doctor_model.addArc(cardiovascular_risk, treatment)

doctor_model.addArc(treatment,risk_of_relapse)


doctor_model.addArc(severity_hypothyroidism, treatment)
doctor_model.addArc(Duration_of_Treatment,TSH_Level )
doctor_model.addArc(TSH_Level,Thyroid_Function_Status )
doctor_model.addArc(treatment, tyroid_eye_disease)
doctor_model.addArc(Eye_Disease_Status, tyroid_eye_disease)
doctor_model.addArc(treatment, Thyroid_Function_Status)
doctor_model.addArc(treatment, Duration_of_Treatment)
doctor_model.addArc(Duration_of_Treatment, Thyroid_Function_Status)

doctor_model.addArc(smoking,tyroid_eye_disease)

doctor_model.addArc(risk_of_relapse, satisfaction)
doctor_model.addArc(Thyroid_Function_Status, satisfaction)
doctor_model.addArc(tyroid_eye_disease, satisfaction)

# CPT for Age Node
doctor_model.cpt(age)[{}] = [0.50, 0.30, 0.20]  # Probabilities for 30-39, 40-49, 50+
# CPT for Smoking Node
doctor_model.cpt(smoking)[{}] = [0.65, 0.35]  # Probabilities for Non-Smoker (0) and Smoker (1)

# CPT for Goiter Size Node based on Age
doctor_model.cpt(goiter_size)[{'Age': 0}] = [0.80, 0.20]  # Age 30-39: Small (80%), Large (20%)
doctor_model.cpt(goiter_size)[{'Age': 1}] = [0.70, 0.30]  # Age 40-49: Small (70%), Large (30%)
doctor_model.cpt(goiter_size)[{'Age': 2}] = [0.60, 0.40]  # Age 50+: Small (60%), Large (40%)


# CPT for General Health Status Node based on Age and Smoking
# General Health Status states: Good (0), Average (1), Poor (2)

# Age: 30-39 (0), Non-smoker (0)
doctor_model.cpt(General_Health_Status)[{'Age': 0, 'Smoking': 0}] = [0.70, 0.20, 0.10]
# Age: 30-39 (0), Smoker (1)
doctor_model.cpt(General_Health_Status)[{'Age': 0, 'Smoking': 1}] = [0.50, 0.30, 0.20]

# Age: 40-49 (1), Non-smoker (0)
doctor_model.cpt(General_Health_Status)[{'Age': 1, 'Smoking': 0}] = [0.50, 0.30, 0.20]
# Age: 40-49 (1), Smoker (1)
doctor_model.cpt(General_Health_Status)[{'Age': 1, 'Smoking': 1}] = [0.30, 0.40, 0.30]

# Age: 50+ (2), Non-smoker (0)
doctor_model.cpt(General_Health_Status)[{'Age': 2, 'Smoking': 0}] = [0.30, 0.40, 0.30]
# Age: 50+ (2), Smoker (1)
doctor_model.cpt(General_Health_Status)[{'Age': 2, 'Smoking': 1}] = [0.10, 0.30, 0.60]

# CPT for Cardiovascular Risk based on Age, Smoking, and General Health Status
# Cardiovascular Risk states: No (0), Yes (1)

# Age 30-39, Non-smoker, Good Health
# CPT for Cardiovascular Risk based on Age, Smoking, and General Health Status
# Cardiovascular Risk states: No (0), Yes (1)

# Age: 30-39 (0), Non-smoker (0)
doctor_model.cpt(cardiovascular_risk)[{'Age': 0, 'Smoking': 0, 'General_Health_Status': 0}] = [0.90, 0.10]  # Good Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 0, 'Smoking': 0, 'General_Health_Status': 1}] = [0.80, 0.20]  # Average Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 0, 'Smoking': 0, 'General_Health_Status': 2}] = [0.60, 0.40]  # Poor Health

# Age: 30-39 (0), Smoker (1)
doctor_model.cpt(cardiovascular_risk)[{'Age': 0, 'Smoking': 1, 'General_Health_Status': 0}] = [0.70, 0.30]  # Good Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 0, 'Smoking': 1, 'General_Health_Status': 1}] = [0.60, 0.40]  # Average Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 0, 'Smoking': 1, 'General_Health_Status': 2}] = [0.40, 0.60]  # Poor Health

# Age: 40-49 (1), Non-smoker (0)
doctor_model.cpt(cardiovascular_risk)[{'Age': 1, 'Smoking': 0, 'General_Health_Status': 0}] = [0.80, 0.20]  # Good Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 1, 'Smoking': 0, 'General_Health_Status': 1}] = [0.70, 0.30]  # Average Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 1, 'Smoking': 0, 'General_Health_Status': 2}] = [0.50, 0.50]  # Poor Health

# Age: 40-49 (1), Smoker (1)
doctor_model.cpt(cardiovascular_risk)[{'Age': 1, 'Smoking': 1, 'General_Health_Status': 0}] = [0.60, 0.40]  # Good Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 1, 'Smoking': 1, 'General_Health_Status': 1}] = [0.50, 0.50]  # Average Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 1, 'Smoking': 1, 'General_Health_Status': 2}] = [0.30, 0.70]  # Poor Health

# Age: 50+ (2), Non-smoker (0)
doctor_model.cpt(cardiovascular_risk)[{'Age': 2, 'Smoking': 0, 'General_Health_Status': 0}] = [0.70, 0.30]  # Good Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 2, 'Smoking': 0, 'General_Health_Status': 1}] = [0.60, 0.40]  # Average Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 2, 'Smoking': 0, 'General_Health_Status': 2}] = [0.40, 0.60]  # Poor Health

# Age: 50+ (2), Smoker (1)
doctor_model.cpt(cardiovascular_risk)[{'Age': 2, 'Smoking': 1, 'General_Health_Status': 0}] = [0.50, 0.50]  # Good Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 2, 'Smoking': 1, 'General_Health_Status': 1}] = [0.40, 0.60]  # Average Health
doctor_model.cpt(cardiovascular_risk)[{'Age': 2, 'Smoking': 1, 'General_Health_Status': 2}] = [0.20, 0.80]  # Poor Health

# CPT for Severity of Hypothyroidism based on Goiter Size
# Goiter Size: Small (0), Large (1)
# Severity of Hypothyroidism: No (0), Yes (1)

# Small goiter
doctor_model.cpt(severity_hypothyroidism)[{'GoiterSize': 0}] = [0.80, 0.20]  # 80% chance of No, 20% chance of Yes

# Large goiter
doctor_model.cpt(severity_hypothyroidism)[{'GoiterSize': 1}] = [0.40, 0.60]  # 40% chance of No, 60% chance of Yes

# CPT for Relapse Risk based on Goiter Size and Treatment Choice
# Relapse Risk states: Low (0), High (1)
# Treatment Choice: ATD (0), RAI (1), Surgery (2)
# Goiter Size: Small (0), Large (1)


doctor_model.cpt(risk_of_relapse)[{'Treatment': 0, 'GoiterSize': 0}] = [0.30, 0.70]  # Increase relapse risk likelihood
doctor_model.cpt(risk_of_relapse)[{'Treatment': 0, 'GoiterSize': 1}] = [0.20, 0.80]


# RAI treatment
doctor_model.cpt(risk_of_relapse)[{'Treatment': 1, 'GoiterSize': 0}] = [0.80, 0.20]  # Small goiter, 20% chance of relapse
doctor_model.cpt(risk_of_relapse)[{'Treatment': 1, 'GoiterSize': 1}] = [0.70, 0.30]  # Large goiter, 30% chance of relapse

# Surgery treatment
doctor_model.cpt(risk_of_relapse)[{'Treatment': 2, 'GoiterSize': 0}] = [0.95, 0.05]  # Small goiter, 5% chance of relapse
doctor_model.cpt(risk_of_relapse)[{'Treatment': 2, 'GoiterSize': 1}] = [0.90, 0.10]  # Large goiter, 10% chance of relapse


# CPT for Eye Disease Status with Three States: Low, Moderate, Severe
# Eye Disease Status: Low (0), Moderate (1), Severe (2)

doctor_model.cpt(Eye_Disease_Status)[{}] = [0.60, 0.25, 0.15]  # Low 60%, Moderate 25%, Severe 15%

# CPT for Thyroid Eye Disease based on Smoking, Eye Disease Status, and Treatment Choice
# Thyroid Eye Disease states: No (0), Yes (1)

# ATD treatment, non-smoker, low initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 0, 'Smoking': 0, 'Eye_Disease_Status': 0}] = [0.90, 0.10]  # Low chance of TED with ATD and no smoking

# ATD treatment, smoker, low initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 0, 'Smoking': 1, 'Eye_Disease_Status': 0}] = [0.70, 0.30]  # Higher chance of TED due to smoking

# ATD treatment, non-smoker, moderate initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 0, 'Smoking': 0, 'Eye_Disease_Status': 1}] = [0.80, 0.20]  # Moderate TED chance

# ATD treatment, smoker, severe initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 0, 'Smoking': 1, 'Eye_Disease_Status': 2}] = [0.50, 0.50]  # Equal chance of persistent or worsening TED

# RAI treatment, non-smoker, low initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 1, 'Smoking': 0, 'Eye_Disease_Status': 0}] = [0.80, 0.20]  # Moderate TED chance due to RAI

# RAI treatment, smoker, low initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 1, 'Smoking': 1, 'Eye_Disease_Status': 0}] = [0.50, 0.50]  # Higher chance of TED due to RAI and smoking

# RAI treatment, smoker, severe initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 1, 'Smoking': 1, 'Eye_Disease_Status': 2}] = [0.20, 0.80]  # Very high chance of persistent or worsening TED

# Surgery treatment, non-smoker, moderate initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 2, 'Smoking': 0, 'Eye_Disease_Status': 1}] = [0.85, 0.15]  # Lower chance of TED due to definitive treatment

# Surgery treatment, smoker, severe initial eye disease
doctor_model.cpt(tyroid_eye_disease)[{'Treatment': 2, 'Smoking': 1, 'Eye_Disease_Status': 2}] = [0.60, 0.40]  # Higher chance of TED but better managed compared to RAI

    
# ATD treatment: most patients require a long treatment duration
doctor_model.cpt(Duration_of_Treatment)[{'Treatment': 0}] = [0.10, 0.20, 0.70]  # Short 10%, Medium 20%, Long 70%

# RAI treatment: typically medium duration but can be short or require follow-up
doctor_model.cpt(Duration_of_Treatment)[{'Treatment': 1}] = [0.30, 0.60, 0.10]  # Short 30%, Medium 60%, Long 10%

# Surgery treatment: shorter active treatment phase but immediate and long-term follow-up
doctor_model.cpt(Duration_of_Treatment)[{'Treatment': 2}] = [0.70, 0.20, 0.10]  # Short 70%, Medium 20%, Long 10%

# Short duration: TSH is more likely to be abnormal (low or high)
doctor_model.cpt(TSH_Level)[{'Duration_of_Treatment': 0}] = [0.40, 0.30, 0.30]  # Low 40%, Normal 30%, High 30%

# Medium duration: TSH is more likely to be normal as the treatment takes effect
doctor_model.cpt(TSH_Level)[{'Duration_of_Treatment': 1}] = [0.20, 0.60, 0.20]  # Low 20%, Normal 60%, High 20%

# Long duration: TSH is most likely normalized, with some risk of overtreatment leading to high TSH
doctor_model.cpt(TSH_Level)[{'Duration_of_Treatment': 2}] = [0.10, 0.70, 0.20]  # Low 10%, Normal 70%, High 20%    

# ATD treatment
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 0, 'Duration_of_Treatment': 0, 'Treatment': 0}] = [0.60, 0.30, 0.10]  # Short duration, low TSH
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 1, 'Duration_of_Treatment': 1, 'Treatment': 0}] = [0.10, 0.80, 0.10]  # Medium duration, normal TSH
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 2, 'Duration_of_Treatment': 2, 'Treatment': 0}] = [0.05, 0.15, 0.80]  # Long duration, high TSH

# RAI treatment
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 0, 'Duration_of_Treatment': 0, 'Treatment': 1}] = [0.50, 0.40, 0.10]  # Short duration, low TSH
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 1, 'Duration_of_Treatment': 1, 'Treatment': 1}] = [0.10, 0.60, 0.30]  # Medium duration, normal TSH
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 2, 'Duration_of_Treatment': 2, 'Treatment': 1}] = [0.05, 0.10, 0.85]  # Long duration, high TSH

# Surgery treatment
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 0, 'Duration_of_Treatment': 0, 'Treatment': 2}] = [0.40, 0.50, 0.10]  # Short duration, low TSH
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 1, 'Duration_of_Treatment': 1, 'Treatment': 2}] = [0.05, 0.30, 0.65]  # Medium duration, normal TSH
doctor_model.cpt(Thyroid_Function_Status)[{'TSH_Level': 2, 'Duration_of_Treatment': 2, 'Treatment': 2}] = [0.02, 0.08, 0.90]  # Long duration, high TSH

#------------------------------------------------------------------------------------------------------------------------------------------------------
# ----------------------------------------------DOCTOR MODEL ENDS -------------------------------------------------------------------------------------
#------------------------------------------------------------------------------------------------------------------------------------------------------


"""These prameters are needed for utility calulation."""
parametersA = np.array(
    [
        ["Thyroid_Function_Status", 1.2, 50], 
        ["Remission", 0.6, 30],  
        ["RelapseRisk", 0.9, 25],  
        ["Tyroid_Eye_Disease", 0.4, 25],  
    ]
)

def calculate_utility_values(model, utility_node_name, parameters):
    """
    This function calculates utilities based on stated weights and increments.
    """
    utility_node_id = model.idFromName(utility_node_name)
    utility_parents = model.parents(utility_node_id)
    parent_names = [model.variable(parent_id).name() for parent_id in utility_parents]

    weights = []
    increments = []

    for parent in parent_names:
        found = False
        for row in parameters:
            if row[0] == parent:
                weights.append(float(row[1]))
                increments.append(float(row[2]))
                found = True
                break

        if not found:
            raise ValueError(f"No parameters found for parent '{parent}'")


    assert len(weights) == len(utility_parents), (
        f"Number of weights must match number of parents. "
        f"Found {len(weights)} weights for {len(utility_parents)} parents."
    )
    assert len(increments) == len(utility_parents), (
        f"Number of increments must match number of parents. "
        f"Found {len(increments)} increments for {len(utility_parents)} parents."
    )

    utilities = {}

    parent_domains = {
        parent_name: model.variable(model.idFromName(parent_name)).labels()
        for parent_name in parent_names
    }

    for parent_state_combination in itertools.product(*parent_domains.values()):
        utility_value = 0

        for i, (parent_name, state) in enumerate(zip(parent_names, parent_state_combination)):
            num_states = len(parent_domains[parent_name])
            state_index = parent_domains[parent_name].index(state)

            incremental_value = (num_states - state_index - 1) * increments[i]
            weighted_value = weights[i] * incremental_value
            utility_value += weighted_value

        utilities[parent_state_combination] = utility_value

    for parent_state_combination, utility_value in utilities.items():
        state_dict = {parent_name: state for parent_name, state in zip(parent_names, parent_state_combination)}
        model.utility(utility_node_id)[state_dict] = utility_value


def find_root_nodes(model):
    """
    This function calculates root chance nodes.
    """
    root_nodes = []
    for node_id in model.nodes():
        if model.isChanceNode(node_id) and not model.parents(node_id):
            root_nodes.append(node_id)
    return root_nodes

def find_all_upstream_nodes(model, node_id):
    """
    This function finds the chance nodes that are upstream of a given node (it will be the decision node)
    """
    
    upstream_nodes = set()
    to_visit = [node_id]

    while to_visit:
        current_node = to_visit.pop()
        for parent in model.parents(current_node):
            if parent not in upstream_nodes:
                upstream_nodes.add(parent)
                to_visit.append(parent)

    return upstream_nodes
    
def VOI(model,node_name,verbose=True):
  def getMEU(evidence):
    limid = gum.ShaferShenoyLIMIDInference(model)
    limid.setEvidence(evidence)
    limid.makeInference()
    meu=limid.MEU()['mean']

    return meu
    
  limid = gum.ShaferShenoyLIMIDInference(model)
  limid.makeInference()
  posterior=limid.posterior(node_name)
  baseline=limid.MEU()['mean']
  
  averageMEU=posterior.expectedValue(getMEU)
  voi=averageMEU-baseline

  return voi
    
def calculate_voi_for_relevant_chance_nodes(model):
    """
    This function calculates VOIs for all chance nodes that are upstream of the decision node.
    """
    total_voi_by_node = {}

  
    decision_node_id = next(node for node in model.nodes() if model.isDecisionNode(node))
    decision_node_name = model.variable(decision_node_id).name()

    relevant_nodes = find_all_upstream_nodes(model, decision_node_id)

    root_nodes = find_root_nodes(model)
    relevant_nodes.update(root_nodes)

    for node_id in relevant_nodes:
        if model.isChanceNode(node_id):
            chance_node_name = model.variable(node_id).name()
            
            try:
                voi = VOI(model, chance_node_name)
                total_voi_by_node[chance_node_name] = voi
            except Exception as e:
                print(f"Error calculating VOI for {chance_node_name}: {e}")

    sorted_voi_by_node = dict(sorted(total_voi_by_node.items(), key=lambda item: item[1], reverse=True))

    print("\nTotal VOI for each relevant chance node (sorted by VOI):")
    for node, voi in sorted_voi_by_node.items():
        print(f"{node}: {voi}")
        
    return sorted_voi_by_node

    
calculate_voi_for_relevant_chance_nodes(doctor_model)