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

# --- Define the network structure ---
# B = Burglary, E = Earthquake, A = Alarm, J = JohnCalls, M = MaryCalls
model = DiscreteBayesianNetwork([
    ("Burglary", "Alarm"),
    ("Earthquake", "Alarm"),
    ("Alarm", "JohnCalls"),
    ("Alarm", "MaryCalls"),
])

# --- Define the CPDs ---
cpd_burglary = TabularCPD(variable="Burglary", variable_card=2,
                          values=[[0.999], [0.001]])  # [False, True]

cpd_earthquake = TabularCPD(variable="Earthquake", variable_card=2,
                            values=[[0.998], [0.002]])

cpd_alarm = TabularCPD(
    variable="Alarm",
    variable_card=2,
    values=[
        [1 - 0.001, 1 - 0.29, 1 - 0.94, 1 - 0.95],  # P(A=False | parents)
        [0.001, 0.29, 0.94, 0.95],                  # P(A=True | parents)
    ],
    evidence=["Burglary", "Earthquake"],
    evidence_card=[2, 2],
)

cpd_johncalls = TabularCPD(
    variable="JohnCalls",
    variable_card=2,
    values=[
        [1 - 0.05, 1 - 0.90],  # P(J=False | A)
        [0.05, 0.90],          # P(J=True | A)
    ],
    evidence=["Alarm"],
    evidence_card=[2],
)

cpd_marycalls = TabularCPD(
    variable="MaryCalls",
    variable_card=2,
    values=[
        [1 - 0.01, 1 - 0.70],
        [0.01, 0.70],
    ],
    evidence=["Alarm"],
    evidence_card=[2],
)

# --- Add CPDs to the model ---
model.add_cpds(cpd_burglary, cpd_earthquake, cpd_alarm, cpd_johncalls, cpd_marycalls)

# --- Check model consistency ---
assert model.check_model()

# --- Perform inference ---
inference = VariableElimination(model)

# Compute posterior P(Burglary | JohnCalls=True, MaryCalls=True)
#posterior = inference.query(variables=["Burglary"],
#                            evidence={"JohnCalls": 1, "MaryCalls": 1})
#print(posterior)

posterior = inference.query(variables=["Burglary"],
                            evidence={"Alarm": 1})
print(posterior)