pip install pgmpy

# <<<CELL>>>

from pgmpy.models import DiscreteBayesianNetwork

# <<<CELL>>>

from pgmpy.factors.discrete import TabularCPD
from pgmpy.inference import VariableElimination
import pandas as pd
import networkx as nx
import matplotlib.pyplot as plt

# <<<CELL>>>

model = DiscreteBayesianNetwork([
    ('A','B'),
    ('C','D'),
    ('B','E'),
    ('B','F'),
    ('D','F'),
    ('F','G'),
    ('E','H'),
    ('E','I'),
])

# <<<CELL>>>

cpt_a = TabularCPD(
    variable='A', variable_card=2,
    values=[[0.9], [0.1]]
)

cpt_c = TabularCPD(
    variable='C', variable_card=2,
    values=[[0.5], [0.5]]
)


cpt_b = TabularCPD(
    variable='B', variable_card=2,
    values=[[0.1, 0.6],
            [0.9, 0.4]],
    evidence=['A'], evidence_card=[2]
)


cpt_d = TabularCPD(
    variable='D', variable_card=2,
    values=[[0.4, 0.7],
            [0.6, 0.3]],
    evidence=['C'], evidence_card=[2]
)


cpt_e = TabularCPD(
    variable='E', variable_card=2,
    values=[[0.8, 0.9],
            [0.2, 0.1]],
    evidence=['B'], evidence_card=[2]
)

cpt_f = TabularCPD(
    variable='F', variable_card=2,
    values=[[1.0, 0.8, 0.6, 0.3],
            [0.0, 0.2, 0.4, 0.7]],
    evidence=['B', 'D'], evidence_card=[2, 2]
)

cpt_g = TabularCPD(
    variable='G', variable_card=2,
    values=[[0.9, 0.1],
            [0.1, 0.9]],
    evidence=['F'], evidence_card=[2]
)

cpt_h = TabularCPD(
    variable='H', variable_card=2,
    values=[[0.3, 0.5],     # H=0
            [0.7, 0.5]],    # H=1
    evidence=['E'], evidence_card=[2]
)

cpt_i = TabularCPD(
    variable='I', variable_card=2,
    values=[[0.3, 0.2],     # I=0
            [0.7, 0.8]],    # I=1
    evidence=['E'], evidence_card=[2]
)

# <<<CELL>>>

model.add_cpds(cpt_a, cpt_b, cpt_c, cpt_d, cpt_e, cpt_f, cpt_g, cpt_h, cpt_i)
print("Model is valid:", model.check_model())

# <<<CELL>>>

for cpd in model.get_cpds():
    print(f"CPT of {cpd.variable}:")
    print(cpd, "\n")

# <<<CELL>>>

nodes = sorted(model.nodes())
adj = pd.DataFrame(0, index=nodes, columns=nodes)
for u, v in model.edges():
    adj.loc[u, v] = 1
print(adj)

# <<<CELL>>>

pos = {'A': (0, 3), 'C': (2, 3), 'B': (0, 2), 'D': (2, 2),
       'E': (-0.6, 1), 'F': (1, 1), 'G': (2.3, 1), 'H': (-1.2, 0), 'I': (0, 0)}
plt.figure(figsize=(7, 6))
nx.draw(nx.DiGraph(model.edges()), pos, with_labels=True, node_size=2000,
        node_color='lightblue', edgecolors='steelblue', font_size=16,
        font_weight='bold', arrowsize=20)
plt.title('Bayesian Network')
plt.show()

# <<<CELL>>>

infer = VariableElimination(model)

print("P(G):")
print(infer.query(variables=['G']))

print("P(G | A=1):")
print(infer.query(variables=['G'], evidence={'A': 1}))

print("P(B | H=1, G=1):")
print(infer.query(variables=['B'], evidence={'H': 1, 'G': 1}))

# <<<CELL>>>

infer = VariableElimination(model)

print("Marginal probabilities P(X = 1):")
for node in sorted(model.nodes()):
    q = infer.query(variables=[node], show_progress=False)
    print(f"  P({node}=1) = {q.values[1]:.4f}")

# <<<CELL>>>

print("P(G | A=1):")
print(infer.query(variables=['G'], evidence={'A': 1}, show_progress=False))

print("P(F | B=1, D=1):   (should equal the CPT value 0.7)")
print(infer.query(variables=['F'], evidence={'B': 1, 'D': 1}, show_progress=False))


print("P(A | G=1):")
print(infer.query(variables=['A'], evidence={'G': 1}, show_progress=False))

print("P(B | H=1, G=1):")
print(infer.query(variables=['B'], evidence={'H': 1, 'G': 1}, show_progress=False))


print("P(D | F=1):")
print(infer.query(variables=['D'], evidence={'F': 1}, show_progress=False))
print("P(D | F=1, B=1):")
print(infer.query(variables=['D'], evidence={'F': 1, 'B': 1}, show_progress=False))


print("Joint P(E, F):")
print(infer.query(variables=['E', 'F'], joint=True, show_progress=False))

# <<<CELL>>>

assignment = {'A': 1, 'B': 1, 'C': 0, 'D': 1, 'E': 0,
              'F': 1, 'G': 1, 'H': 1, 'I': 1}
p = 1.0
for cpd in model.get_cpds():
    var = cpd.variable
    idx = [assignment[var]] + [assignment[e] for e in cpd.variables[1:]]
    p *= cpd.values[tuple(idx)]
print(f"P(full assignment) = {p:.6f}")


print("MAP of A, B, C, D given G=1, H=1:")
print(infer.map_query(variables=['A', 'B', 'C', 'D'],
                      evidence={'G': 1, 'H': 1}, show_progress=False))


print("Is A independent of D?          ", not model.is_dconnected('A', 'D'))
print("Is A independent of D given F?  ", not model.is_dconnected('A', 'D', observed=['F']))
print("Is H independent of I given E?  ", not model.is_dconnected('H', 'I', observed=['E']))
print("Is A independent of H given B?  ", not model.is_dconnected('A', 'H', observed=['B']))

# <<<CELL>>>

# %%
def prob(query):
    """
    Evaluate a probability query written in textbook notation.
      lowercase letter  -> variable is True  (1)   e.g.  g
      ~ before letter   -> variable is False (0)   e.g.  ~g
      '|' separates the query part from the evidence (given) part
      ',' separates variables
    Examples:  "g"   "g | a"   "g, ~i | a, c"   "b | h, g, ~d"
    """
    query = query.replace(" ", "").replace("P(", "").rstrip(")")
    left, _, right = query.partition("|")

    def parse(part):
        out = {}
        for term in filter(None, part.split(",")):
            if term.startswith("~"):
                out[term[1:].upper()] = 0
            else:
                out[term.upper()] = 1
        return out

    targets, evidence = parse(left), parse(right)

    result = infer.query(variables=list(targets), evidence=evidence or None,
                         joint=True, show_progress=False)
    return result.get_value(**targets)


from itertools import product

def joint_prob(assign):
    p = 1.0
    for cpd in model.get_cpds():
        idx = [assign[cpd.variable]] + [assign[e] for e in cpd.variables[1:]]
        p *= cpd.values[tuple(idx)]
    return p

def prob_enum(targets, evidence):
    nodes = list(model.nodes())
    num = den = 0.0
    for values in product([0, 1], repeat=len(nodes)):
        w = dict(zip(nodes, values))
        if all(w[k] == v for k, v in evidence.items()):
            p = joint_prob(w)
            den += p
            if all(w[k] == v for k, v in targets.items()):
                num += p
    return num / den


queries = [
    "P(g)",
    "P(g | a)",
    "P(g | ~a)",
    "P(g | c)",
    "P(g | a, c)",
    "P(g | b, ~d)",
    "P(a | g)",
    "P(a | g, h)",
    "P(b | h, g, ~i)",
    "P(d | f)",
    "P(d | f, b)",
    "P(g, h | a)",
    "P(g, ~i | a, c)",
    "P(f, g | ~b, d)",
    "P(a, c | g, ~h, i)",
]

print(f"{'Query':<25}{'pgmpy':>10}{'enumeration':>14}")
print("-" * 49)
for q in queries:
    body = q.replace(" ", "")[2:-1]
    left, _, right = body.partition("|")
    parse = lambda s: {t.lstrip('~').upper(): 0 if t.startswith('~') else 1
                       for t in filter(None, s.split(","))}
    print(f"{q:<25}{prob(q):>10.4f}{prob_enum(parse(left), parse(right)):>14.4f}")

# %%
# Ask your own query here:
print(prob("P(g | a, ~h)"))
