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) +"\n"
            print(s)
        elif len(parents) == 1:
            p = parents[0]
            s += str(yes_prob[0]) + " :: " + str(node) + " <- " + str(p) + "\n"
            s += str(yes_prob[1]) + " :: " + str(node) + " <- -" + str(p) + "\n"
            print(s)
        else:
            yes1 = yes_prob[0]
            yes2 = yes_prob[1]
            p0 = parents[0]
            p1 = parents[1]
            s += str(yes1[0]) + " :: " + str(node) + " <- " + str(p0) + ", " + str(p1) + "\n"
            s += str(yes2[0]) + " :: " + str(node) + " <- -" + str(p0) + ", " + str(p1) + "\n"
            s += str(yes1[1]) + " :: " + str(node) + " <- " + str(p0) + ", -" + str(p1) + "\n"
            s += str(yes2[1]) + " :: " + str(node) + " <- -" + str(p0) + ", -" + str(p1) + "\n"
            print(s)