<filename>scripts/student_pgm.py
#https://github.com/pgmpy/pgmpy_notebook/blob/master/notebooks/2.%20Bayesian%20Networks.ipynb

from pgmpy.models import BayesianModel
from pgmpy.factors.discrete import TabularCPD

# Defining the model structure. We can define the network by just passing a list of edges.
model = BayesianModel([('D', 'G'), ('I', 'G'), ('G', 'L'), ('I', 'S')])

# Defining individual CPDs.
cpd_d = TabularCPD(variable='D', variable_card=2, values=[[0.6, 0.4]])
cpd_i = TabularCPD(variable='I', variable_card=2, values=[[0.7, 0.3]])

# The representation of CPD in pgmpy is a bit different than the CPD shown in the above picture. In pgmpy the colums
# are the evidences and rows are the states of the variable. So the grade CPD is represented like this:
#
#    +---------+---------+---------+---------+---------+
#    | diff    | intel_0 | intel_0 | intel_1 | intel_1 |
#    +---------+---------+---------+---------+---------+
#    | intel   | diff_0  | diff_1  | diff_0  | diff_1  |
#    +---------+---------+---------+---------+---------+
#    | grade_0 | 0.3     | 0.05    | 0.9     | 0.5     |
#    +---------+---------+---------+---------+---------+
#    | grade_1 | 0.4     | 0.25    | 0.08    | 0.3     |
#    +---------+---------+---------+---------+---------+
#    | grade_2 | 0.3     | 0.7     | 0.02    | 0.2     |
#    +---------+---------+---------+---------+---------+

cpd_g = TabularCPD(variable='G', variable_card=3, 
                   values=[[0.3, 0.05, 0.9,  0.5],
                           [0.4, 0.25, 0.08, 0.3],
                           [0.3, 0.7,  0.02, 0.2]],
                  evidence=['I', 'D'],
                  evidence_card=[2, 2])

cpd_l = TabularCPD(variable='L', variable_card=2, 
                   values=[[0.1, 0.4, 0.99],
                           [0.9, 0.6, 0.01]],
                   evidence=['G'],
                   evidence_card=[3])

cpd_s = TabularCPD(variable='S', variable_card=2,
                   values=[[0.95, 0.2],
                           [0.05, 0.8]],
                   evidence=['I'],
                   evidence_card=[2])

# Associating the CPDs with the network
model.add_cpds(cpd_d, cpd_i, cpd_g, cpd_l, cpd_s)

# check_model checks for the network structure and CPDs and verifies that the CPDs are correctly 
# defined and sum to 1.
model.check_model()

from pgmpy.inference import VariableElimination
infer = VariableElimination(model)

# p(I=1)=0.3
print(infer.query(['I']) ['I'])




# P(I=1|G=0) = 0.6133
print(infer.query(['I'], evidence={'G': 0}) ['I'])

# P(I=1|G=0,D=0) = 0.5625
print(infer.query(['I'], evidence={'G': 0, 'D': 0}) ['I'])

# P(S=1|G=0) = 0.5099
print(infer.query(['S'], evidence={'G': 0}) ['S'])

# P(S=1|G=0,D=0) = 0.4719
print(infer.query(['S'], evidence={'G': 0, 'D': 0}) ['S'])



# P(I=1|L=0) = 0.14
print(infer.query(['I'], evidence={'L': 0}) ['I'])

# P(I=1|L=0,D=1) = 0.16
print(infer.query(['I'], evidence={'L': 0, 'D': 1}) ['I'])

# P(I=1|L=0,D=1,S=1) = 0.7597
print(infer.query(['I'], evidence={'L': 0, 'D': 1, 'S': 1}) ['I'])

# P(I=1|L=0,S=1) = 0.7226
print(infer.query(['I'], evidence={'L': 0,  'S': 1}) ['I'])
