"""Script to extract BNs from maxwalksat output
"""

import sys
import cPickle
from gPy.Structures import ADG
from Tkinter import *

bn_readin = False

probadg = cPickle.load(open(sys.argv[1]))
store = probadg._wcnf_info['atom_store']
if 'families' in store:
    dkt = store['families']
    def alter_bn(bn,(parents,child)):
        bn.put_family(child,parents)
elif 'parents' in store:
    dkt = store['parents']
    def alter_bn(bn,(parent,child)):
        bn.put_arrow(parent,child)
else:
    raise ValueError('Neither a family nor a parent encoding was used, something is wrong')

decode = {}    
for key, atom in dkt.items():
    decode[atom] = key

for line in open((sys.argv[2])):
    if line.startswith('Begin assign'):
        bn_readin = True
        bn = ADG(probadg._variables)
    elif line.startswith('End assign'):
        bn_readin = False
        print 'Bayesian network:'
        print bn
        print
        print 'Essential graph:'
        ess = bn.essential_graph()
        print ess
        print
        root = Tk()
        bn.gui_display(root)
        ess.gui_display(root)
        root.mainloop()
    elif bn_readin:
        for atom in line.split():
            try:
                alter_bn(bn,decode[int(atom)])
            except KeyError:
                pass