#https://pgmpy.org/examples/Monty%20Hall%20Problem.html

from pgmpy.models import DiscreteBayesianNetwork
from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import VariableElimination

# Define the structure of the Bayesian Network
model = DiscreteBayesianNetwork([('guest', 'monty'), ('prize', 'monty')]) # guest & proze = parents, monty = child

# Define the Conditional Probability Distribution (CPD) for the guest's choice (completely random)
cpd_guest = TabularCPD(variable='guest', 
                       variable_card=3, 
                       values=[[1/3], [1/3], [1/3]], 
                       state_names={'guest': ['A', 'B', 'C']})

# Define the CPD for the prize location (also random)
cpd_prize = TabularCPD(variable='prize', 
                       variable_card=3, 
                       values=[[1/3], [1/3], [1/3]], 
                       state_names={'prize': ['A', 'B', 'C']})

# Define the CPD for Monty's choice, conditional on guest and prize
cpd_monty = TabularCPD(
    variable='monty', 
    variable_card=3,
    values=[
        # monty = A
        [0.0, 0.0, 0.0, 0.0, 0.5, 1.0, 0.0, 1.0, 0.5],
        # monty = B
        [0.5, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.5],
        # monty = C
        [0.5, 1.0, 0.0, 1.0, 0.5, 0.0, 0.0, 0.0, 0.0],
    ],
    evidence=['guest', 'prize'], #nodes that influence "monty"
    evidence_card=[3, 3], #each parent has 3 states (A,B,C)
    state_names={'guest': ['A','B','C'], 'prize': ['A','B','C'], 'monty': ['A','B','C']}
)
 
#guest	prize	monty=A	    monty=B	    monty=C   (0.0: impossible, 0.5:random choice between two, 1.0:certain)
#A	     A	     0.0	      0.5	      0.5
#A	     B	     0.0	      0.0	      1.0
#A	     C	     0.0	      1.0	      0.0
#B	     A	     0.0	      0.0	      1.0
#B	     B	     0.5	      0.0	      0.5
#B	     C	     1.0	      0.0	      0.0
#C	     A	     0.0	      1.0	      0.0
#C	     B	     1.0	      0.0	      0.0
#C	     C	     0.5	      0.5	      0.0

# Add CPDs to the model
model.add_cpds(cpd_guest, cpd_prize, cpd_monty)

# Verify there are no errors in structure, CPDs, etc
assert model.check_model() #works only for bayesian networks and markov models (found in pgmpy library)

# Calculate probabilities that the prize is behind an X door
inference = VariableElimination(model) #Create inference engine to perform exact Bayesian reasoning (apply Bayes' rule)

# Query the network for the first scenario (guest chooses 'A')
result = inference.query(variables=['prize'], evidence={'guest': 'A'}) #update belief: where’s the prize now given player's and Monty's choices?
print("Beliefs for 'prize' given 'guest': 'A'")
print(result)

# Query the network for the second scenario (guest chooses 'A', Monty chooses 'B')
result = inference.query(variables=['prize'], evidence={'guest': 'A', 'monty': 'B'}) #update belief: where’s the prize now that the player chose A and Monty opened B?
print("\nBeliefs for 'prize' given 'guest': 'A' and 'monty': 'B'")
print(result)