"""
=====================
Bayesian Networs
=====================
"""

import bnlearn as bn
from utils import prepare_data, print_version_info

# %%

print_version_info()

# %%

data, *_ = prepare_data()

# %%

# Structure learning
# model = bn.structure_learning.fit(dfnum, methodtype='cl',
# black_list=['Embarked','Parch','Name'], root_node='Survived', bw_list_method='nodes')
model = bn.structure_learning.fit(data)

# Plot
G = bn.plot(model, interactive=False)

# %%
# Compute edge strength with the chi_square test statistic
model = bn.independence_test(model, data, test='chi_square', prune=True)

# %%

# Plot
bn.plot(model, interactive=False, pos=G['pos'])

# %%

# Parameter learning
model = bn.parameter_learning.fit(model, data)

# %%

# Make inference
query = bn.inference.fit(
    model,
    variables=['Catalyst type'],
    evidence={'solution pH':3.0,
              'Ci (mg/L)':5.0})

# %%

print(query)