from pgmpy.readwrite import BIFReader
import networkx as nx

import pydot
from networkx.drawing.nx_pydot import graphviz_layout
from bninput import *

def get_yes_probabilities(model):
    nodes = model.nodes()

    yes_probabilities = {}

    for node in nodes:
        cpd = model.get_cpds(node)
        state_names = cpd.state_names[node]
        yes_index = state_names.index("yes")
        yes_prob = cpd.values[yes_index]
        yes_probabilities[node] = yes_prob

    return yes_probabilities

if __name__ == "__main__":
    summary = summary_dag("asia2.bif")

    model = summary["bnmodel"]
    probabilities = get_yes_probabilities(model)
    for node, yes_prob in probabilities.items():
        parents = model.get_parents(node)
        s = ""
        if len(parents) == 0:
            s += str(yes_prob) + " :: " + str(node)
            print(s)
        elif len(parents) == 1:

            s += str(yes_prob[0]) + " :: " + str(node) + " <- " + str(parents) + "\n"

            s += str(yes_prob[1]) + " :: " + str(node) + " <- -" + str(parents) + "\n"

            print(s)
        else:
            print("Node:", node)
            print("Probability of 'yes':", yes_prob)

            print()