#Project 1 for the University of Tulsa's CS-7313 Adv. AI Course
#Approximate Inference Methods for Bayesian Networks
#Professor: Dr. Sen, Fall 2021
#Noah Schrick - 1492657

import json
import random
import math
import sys
import os

#pwd = os.getcwd()
#sys.path.append(pwd +'/gen-bn/')
#sys.path.append("./gen-bn")

import gen_bn.gen_bn

#from collections import defaultdict

def main():
    #Generate a new BN. Specify type and number of nodes in network
    gen_json("dag", 5)
    
    #Get our BN
    bayes_net = import_bayes()
    
    #Generate random evidence
    E = gen_ev(bayes_net)
    
    #Get W from LW
    W = likelihood_weighting(1, E, bayes_net, 10)
    
    #Print if desired
    print()
    for key, value in W.items():
        print(key, ' : ', value)

#Generate a new BN.
#Input: Type ("dag", or "polytree")
#Input: Number of nodes in the network
def gen_json(type, num_nodes):
    os.chdir("./gen_bn")
    os.system('python gen_bn.py' + ' ' + type + ' ' + str(num_nodes))
 
#Import the BN from the json
def import_bayes():
    with open ("gen_bn/bn.json") as json_file:
        bayes_json = json.load(json_file)
    json_file.close()
    return bayes_json

#Generate a random set of evidence
def gen_ev(bayes_net):
    total_nodes = len(bayes_net)
    #Arbitrarily, let's only generate total_nodes/2 (rounded up) evidence variables at most, but at least 1
    num_ev = random.randint(1, int(math.ceil(total_nodes/2)))
    fixed_ev = []
    #Go through and generate nodes that will be fixed
    for i in range(num_ev):
        fixed_var = random.randint(0, total_nodes-1)
        if fixed_var not in fixed_ev:
            fixed_ev.append(fixed_var)
    #Now generate random values for the ev
    #Randomly generate a double. >=0.5 will be "True", <0.5 will be "False"
    E = {}
    for i in fixed_ev:
        val_p = random.random()
        if val_p >= 0.5:
            E[str(i)] = True
        else:
            E[str(i)] = False
    
    return E


#Checks if node has parents
def is_root(node, BN):
    return (BN[node]["parents"]) == []

#Return a list of the root nodes
def get_root(BN):
    roots = []
    for node in range(len(BN)):
        if ((BN[str(node)]["parents"]) == []):
            roots.append(str(node))
    return roots

#Get parents of a node
def get_parents(node, BN):
    return BN[str(node)]["parents"]

"""NOTES"""
#(bayes_json["x"]): the information about node x (an int)
#(bayes_json["x"]["parents"] the information about node x's parents
#bayes_json["x"]["prob"][0][0] returns the first set of truth table (0, 0), where [1][0] is the second (0,1)
#bayes_json["x"]["prob"][parent][1] returns probability of the set evidence variable
#E is a dict in the form of {"Node" : Value}

#Compute the estimate of P(X|e), where X is the query variable, and e is the observed value for variables E
def likelihood_weighting(X, e, bayes_net, num_samples):
    W = {}
    T = 0
    F = 0

    for i in range(num_samples):
        w = 1
        #Holds all the info on the samples. EX: ~b, ~e, a, ~j, m
        samples = {}

        #Get all the roots to save traversion costs
        root_nodes = get_root(bayes_net)
        #Go through all the roots to get probabilities
        for root in root_nodes:
            #If the root is an evidence variables
            if root in e and root not in samples:
                #Just set the value to the observed value
                samples[root] = e[root]
                #Adjust the weight accordingly
                w = w * bayes_net[root]["prob"][0][1]
            else:
                #Otherwise, sample randomly
                if root not in samples:
                    rand_prob = random.random()
                    if rand_prob >= bayes_net[root]["prob"][0][1]:
                        samples[root] = True
                    else:
                        samples[root] = False

        #Now go through the BN for non-root nodes
        for node in bayes_net:
            if node not in samples:
                #Get the probability, updated sample dict, and weight
                samples, prob, w = get_probability(str(node), samples, e, bayes_net, w)

        #We now need to write to W
        #If this sample is already in W, don't add a new sample - only adjust the weight
        written = False
        for tmp in range(len(W)):
            #If sample is in W
            if samples in W[tmp].values():
                #Pull the weight that's associated with the sample
                key = list(W[tmp].items())[0][0]
                #Add the new weight to the existing weight  
                new_key = key + w
                #Store it all back into W
                W[tmp] = {new_key : samples}
                #Make note that we've already written to W in this loop, so we don't write it again
                written = True

        #If the sample wasn't already in W, put it in there now    
        if not written:                 
            W[len(W)] = {w : samples}
    return W


#Return the probability of a node and the value dict, given the current evidence and fixed evidence
#Uses recursion to pull probabilites and values down through the network
def get_probability(node, samples, ev, bayes_net, w):
    parents = get_parents(node, bayes_net)
    for parent in parents:
        #If we already know the value of the parent, no need to reobtain
        if str(parent) in samples:
            continue
        #If we don't know the value, then we need to get it
        else:
            gparents = get_parents(parent, bayes_net)
            #If we have all of parent's parents, then we can just get the probability
            if all(eles in samples for eles in gparents):
                samples, prob,  w = translate_ev(gparents, parent, ev, samples, bayes_net, w)

            #Otherwise, we need to get values for the ancestor nodes - use recursion
            else:
                for gparent in gparents:
                    if gparent not in samples:
                        get_probability(gparent, samples, ev, bayes_net, w)        
                samples, prob,  w = translate_ev(gparents, parent, ev, samples, bayes_net, w)
    
    #Now that we have all the parents' values, we can get the node value, probability, and update samples
    samples, prob, w = translate_ev(parents, node, ev, samples, bayes_net, w)
    return samples, prob, w

#Given a node and its parents, determine the node's value and it's probability
def translate_ev(parents, node, ev, samples, bayes_net, w):
    #Sort in ascending order
    parents.sort()
    node = str(node)
    value_list = []
    for parent in parents: 
        value = samples[str(parent)]
        value_list.append(value)
    
    #See if this is an evidence node
    if node in ev:
        samples[node] = ev[node]
        get_weight = True
    else:
        get_weight = False
    #The truth table has 2^parents entries
    for i in range(2**len(parents)):
        #If the truth table matches the value combination we have
        if bayes_net[str(node)]["prob"][i][0] == value_list:
            #Sample randomly
            rand_prob = random.random()
            table_prob = bayes_net[str(node)]["prob"][i][1]
            if rand_prob >= table_prob:
                    samples[str(node)] = True
            else:
                samples[str(node)] = False
                table_prob = 1-table_prob
    if(get_weight):
        w = w * table_prob

    return samples, table_prob, w
      

def gibbs_sampling():
    print("Hello")

def metropolis_hastings():
    print("Hello")

if __name__ == '__main__':
    main()