##Three sample bayes nets are defined here. ##
from bnetbase import *

VisitAsia = Variable('Visit_To_Asia', ['visit', 'no-visit'])
F1 = Factor("F1", [VisitAsia])
F1.add_values([['visit', 0.01], ['no-visit', 0.99]])

Smoking = Variable('Smoking', ['smoker', 'non-smoker'])
F2 = Factor("F2", [Smoking])
F2.add_values([['smoker', 0.5], ['non-smoker', 0.5]])

Tuberculosis = Variable('Tuberculosis', ['present', 'absent'])
F3 = Factor("F3", [Tuberculosis, VisitAsia])
F3.add_values([['present', 'visit', 0.05],
               ['present', 'no-visit', 0.01],
               ['absent', 'visit', 0.95],
               ['absent', 'no-visit', 0.99]])

Cancer = Variable('Lung Cancer', ['present', 'absent'])
F4 = Factor("F4", [Cancer, Smoking])
F4.add_values([['present', 'smoker', 0.10],
               ['present', 'non-smoker', 0.01],
               ['absent', 'smoker', 0.90],
               ['absent', 'non-smoker', 0.99]])

Bronchitis = Variable('Bronchitis', ['present', 'absent'])
F5 = Factor("F5", [Bronchitis, Smoking])
F5.add_values([['present', 'smoker', 0.60],
               ['present', 'non-smoker', 0.30],
               ['absent', 'smoker', 0.40],
               ['absent', 'non-smoker', 0.70]])

TBorCA = Variable('Tuberculosis or Lung Cancer', ['true', 'false'])
F6 = Factor("F6", [TBorCA, Tuberculosis, Cancer])
F6.add_values([['true', 'present', 'present', 1.0],
               ['true', 'present', 'absent', 1.0],
               ['true', 'absent', 'present', 1.0],
               ['true', 'absent', 'absent', 0],
               ['false', 'present', 'present', 0],
               ['false', 'present', 'absent', 0],
               ['false', 'absent', 'present', 0],
               ['false', 'absent', 'absent', 1]])


Dyspnea = Variable('Dyspnea', ['present', 'absent'])
F7 = Factor("F7", [Dyspnea, TBorCA, Bronchitis])
F7.add_values([['present', 'true', 'present', 0.9],
               ['present', 'true', 'absent', 0.7],
               ['present', 'false', 'present', 0.8],
               ['present', 'false', 'absent', 0.1],
               ['absent', 'true', 'present', 0.1],
               ['absent', 'true', 'absent', 0.3],
               ['absent', 'false', 'present', 0.2],
               ['absent', 'false', 'absent', 0.9]])


Xray = Variable('XRay Result', ['abnormal', 'normal'])
F8 = Factor("F8", [Xray, TBorCA])
F8.add_values([['abnormal', 'true', 0.98],
               ['abnormal', 'false', 0.05],
               ['normal', 'true', 0.02],
               ['normal', 'false', 0.95]])

Asia = BN("Asia", [VisitAsia, Smoking, Tuberculosis, Cancer,
                   Bronchitis, TBorCA, Dyspnea, Xray],
                   [F1, F2, F3, F4, F5, F6, F7, F8])

## E,B,S,w,G example from sample questions
'''E = Variable('E', ['e', '-e'])
B = Variable('B', ['b', '-b'])
S = Variable('S', ['s', '-s'])
G = Variable('G', ['g', '-g'])
W = Variable('W', ['w', '-w'])
FE = Factor('P(E)', [E])
FB = Factor('P(B)', [B])
FS = Factor('P(S|E,B)', [S, E, B])
FG = Factor('P(G|S)', [G,S])
FW = Factor('P(W|S)', [W,S])

FE.add_values([['e',0.1], ['-e', 0.9]])
FB.add_values([['b', 0.1], ['-b', 0.9]])
FS.add_values([['s', 'e', 'b', .9], ['s', 'e', '-b', .2], ['s', '-e', 'b', .8],['s', '-e', '-b', 0],
               ['-s', 'e', 'b', .1], ['-s', 'e', '-b', .8], ['-s', '-e', 'b', .2],['-s', '-e', '-b', 1]])
FG.add_values([['g', 's', 0.5], ['g', '-s', 0], ['-g', 's', 0.5], ['-g', '-s', 1]])
FW.add_values([['w', 's', 0.8], ['w', '-s', .2], ['-w', 's', 0.2], ['-w', '-s', 0.8]])

Q3 = BN('SampleQ4', [E,B,S,G,W], [FE,FB,FS,FG,FW])
'''

A = Variable('A', ['a', '-a'])
B = Variable('B', ['b', '-b'])
C = Variable('C', ['c', '-c'])
D = Variable('D', ['d', '-d'])
E = Variable('E', ['e', '-e'])
F = Variable('F', ['f', '-f'])
G = Variable('G', ['g', '-g'])
H = Variable('H', ['h', '-h'])
I = Variable('I', ['i', '-i'])

FA = Factor('P(A)', [A])
FB = Factor('P(B|A,H)', [B, A, H])
FC = Factor('P(C|B,G)', [C, B, G])
FD = Factor('P(D|C,F)', [D, C, F])
FE = Factor('P(E|C)', [E,C])
FF = Factor('P(F)', [F])
FG = Factor('P(G)', [G])
FH = Factor('P(H)', [H])
FI = Factor('P(I|B)', [I,B])

FA.add_values([['a',0.9], ['-a', 0.1]])
FB.add_values([['b', 'a', 'h', 1], ['b', 'a', '-h', 0], ['b', '-a', 'h', 0.5], ['b', '-a', '-h', 0.6],
               ['-b', 'a', 'h', 0], ['-b', 'a', '-h', 1], ['-b', '-a', 'h', 0.5], ['-b', '-a', '-h', 0.4]])
FC.add_values([['c', 'b', 'g', .9], ['c', 'b', '-g', .9], ['c', '-b', 'g', .1],['c', '-b', '-g', 1],
               ['-c', 'b', 'g', .1], ['-c', 'b', '-g', .1], ['-c', '-b', 'g', .9],['-c', '-b', '-g', 0]])
FD.add_values([['d', 'c', 'f', 0], ['d', 'c', '-f', 1], ['d', '-c', 'f', .7],['d', '-c', '-f', .2],
               ['-d', 'c', 'f', 1], ['-d', 'c', '-f', 0], ['-d', '-c', 'f', .3],['-d', '-c', '-f', .8]])
FE.add_values([['e', 'c', .2], ['e', '-c', .4], ['-e', 'c', 0.8], ['-e', '-c', .6]])
FF.add_values([['f',0.1], ['-f', 0.9]])
FG.add_values([['g',1], ['-g', 0]])
FH.add_values([['h',0.5], ['-h', 0.5]])
FI.add_values([['i', 'b', .3], ['i', '-b', .9], ['-i', 'b', 0.7], ['-i', '-b', .1]])

Q3 = BN('SampleQ4', [A,B,C,D,E,F,G,H,I], [FA,FB,FC,FD,FE,FF,FG,FH,FI])


if __name__ == '__main__':
    #(a)
    
    print("Test 1 ....", end = '')
    B.set_evidence('b')
    probs = VE(Q3, I, [B])
    '''if probs[0] == 1 and probs[1] == 0:
      print("passed.")
    else:
      print("failed.")'''
    
    print('P(i|b) = {} P(-i|b) = {}'.format(probs[0], probs[1]))

    #(b)
    print("Test 2 ....", end = '')
    H.set_evidence('h')
    B.set_evidence('b')
    probs = VE(Q3, I, [B,H])
    '''if abs(probs[0] - 0.68) < 0.0001 and abs(probs[1] - 0.32) < 0.0001:
      print("passed.")
    else:
      print("failed.")    '''
    
    print('P(i|bh) = {} P(-i|bh) = {}'.format(probs[0],probs[1]))

    #(c)
    print("Test 3 ....", end = '')
    A.set_evidence('a')
    B.set_evidence('-b')
    probs1 = VE(Q3, H, [A,B])
    #S.set_evidence('-s')
    #probs2 = VE(Q3, G, [S])
    '''if probs1[0] == 0.5 and probs1[1] == 0.5 and probs2[0] == 0.0 and probs2[1] == 1.0:
      print("passed.")
    else:
      print("failed.")  '''
    print('P(h|a-b) = {} P(-h|a-b) = {}'.format(probs1[0],probs1[1]))

    #(d)
    print("Test 4 ....", end = '')
    B.set_evidence('-b')
    probs1 = VE(Q3, H, [B])
    #S.set_evidence('-s')
    #probs2 = VE(Q3, G, [S])
    '''if probs1[0] == 0.5 and probs1[1] == 0.5 and probs2[0] == 0.0 and probs2[1] == 1.0:
      print("passed.")
    else:
      print("failed.")  '''
    print('P(h|-b) = {} P(-h|-b) = {}'.format(probs1[0],probs1[1]))
'''
    print("Test 5 ....", end = '')
    S.set_evidence('-s')
    W.set_evidence('w')
    probs3 = VE(Q3, G, [S,W])
    S.set_evidence('-s')
    W.set_evidence('-w')
    probs4 = VE(Q3, G, [S,W])
    if probs3[0] == 0.0 and probs3[1] == 1.0 and probs4[0] == 0.0 and probs4[1] == 1.0:
      print("passed.")
    else:
      print("failed.") 
    print('P(g|-s,w) = {} P(-g|-s,w) = {} P(g|-s,-w) = {} P(-g|-s,-w) = {}'.format(probs3[0],probs3[1],probs4[0],probs4[1]))

    #(f) 
    print("Test 6 ....", end = '')
    W.set_evidence('w')
    probs1 = VE(Q3, G, [W])
    W.set_evidence('-w')
    probs2 = VE(Q3, G, [W])
    if abs(probs1[0] - 0.15265998457979954) < 0.0001 and abs(probs1[1] - 0.8473400154202004) < 0.0001 and abs(probs2[0] - 0.01336753983256819) < 0.0001 and abs(probs2[1] - 0.9866324601674318) < 0.0001:
      print("passed.")
    else:
      print("failed.")      
    print('P(g|w) = {} P(-g|w) = {} P(g|-w) = {} P(-g|-w) = {}'.format(probs1[0],probs1[1],probs2[0],probs2[1]))

    #(h)
    print("Test 7 ....", end = '')
    probs = VE(Q3, G, [])
    if abs(probs[0] - 0.04950000000000001) < .0001 and abs(probs[1] - 0.9505) < .0001:
      print("passed.")
    else:
      print("failed.")      
    print('P(g) = {} P(-g) = {}'.format(probs[0], probs[1]))
    
    print("Test 8 ....", end = '')    
    probs = VE(Q3, E, [])
    if abs(probs[0] - 0.1) < 0.0001 and abs(probs[1] - 0.9) < 0.0001:
      print("passed.")
    else:
      print("failed.")      
    print('P(e) = {} P(-e) = {}'.format(probs[0], probs[1]))

    ## NEW TESTS ##
    print("\nMultiply Factors Tests")
    print("Test 1 ....", end = '')    
    factor = multiply_factors([FE])
    values = (factor.get_value(['e']), factor.get_value(['-e']))
    if values[0] == 0.1 and values[1] == 0.9:
      print("passed.")
    else:
      print("failed.")      
    print('P(e) = {} P(-e) = {}'.format(values[0], values[1]))

    print("Test 2 ....", end = '')    
    factor = multiply_factors([FB, FE])
    tests = []
    values = []
    for e_val in E.domain():
      for b_val in B.domain():
        try:
          value = factor.get_value([e_val, b_val])
          values.append(value)
        except ValueError:
          value = factor.get_value([b_val, e_val])
          values.append(value)
        tests.append(value == FE.get_value([e_val])*FB.get_value([b_val]))
    if all(tests):
      print("passed.")
    else:
      print("failed.")      
    print('P(e,b) = {} P(-e,b) = {} P(e,-b) = {} P(-e,-b) = {}'.format(values[0], values[1], values[2], values[3]))

    print("Test 3 ....", end = '')    
    factor = multiply_factors([FE, FS])
    tests = []
    values = []
    for e_val in E.domain():
      for b_val in B.domain():
        for s_val in S.domain():
          try:
            value = factor.get_value([e_val, b_val, s_val])
            values.append(value)
          except ValueError:
            try:
              value = factor.get_value([s_val, b_val, e_val])
              values.append(value)
            except ValueError:
              try:
                value = factor.get_value([b_val, s_val, e_val])
                values.append(value)
              except ValueError:
                try:
                  value = factor.get_value([s_val, e_val, b_val])
                  values.append(value)
                except ValueError:
                  try:
                    value = factor.get_value([b_val, e_val, s_val])
                    values.append(value)
                  except ValueError:
                    value = factor.get_value([e_val, s_val, b_val])
                    values.append(value)
        tests.append(value == FE.get_value([e_val])*FS.get_value([s_val, e_val, b_val]))
    if all(tests):
      print("passed.")
    else:
      print("failed.")      
    print('P(e,s,b) = {} P(-e,s,b) = {} P(e,-s,-b) = {} P(-e,s,-b) = {}'.format(values[0], values[1], values[2], values[3]))
    

    ##
    print("\nSum Out Variable Tests")
    print("Test 1 ....", end = '')    
    factor = sum_out_variable(FE, E)
    value = factor.get_value_at_current_assignments()
    if value == 1:
      print("passed.")
    else:
      print("failed.")
    print('sum_e P(e) = {}'.format(value))
    
    print("Test 2 ....", end = '')    
    factor = sum_out_variable(FS, E)
    values = (factor.get_value(["s", "b"]), factor.get_value(["s", "-b"]), factor.get_value(["-s", "b"]), factor.get_value(["-s", "-b"]))
    tests = (abs(values[0] - 1.7) < 0.001, abs(values[1] - 0.2) < 0.001, abs(values[2] - 0.3) < 0.001, abs(values[3] - 1.8) < 0.001)
    if all(tests):
      print("passed.")
    else:
      print("failed.")
    print('P(S = s | B = b) = {} P(S = s | B = -b) = {} P(S = -s | B = b) = {} P(S = -s | B = -b) = {}'.format(values[0], values[1], values[2], values[3]))
    
    ##
    print("\nRestrict Factor Tests")
    print("Test 1 ....", end = '')    
    factor = restrict_factor(FE, E, 'e')
    value = factor.get_value_at_current_assignments()
    if value == 0.1:
      print("passed.")
    else:
      print("failed.")
    print('P(E=e) = {}'.format(value))

    print("Test 2 ....", end = '')    
    factor = restrict_factor(FG, S, 's')
    value = factor.get_value_at_current_assignments()
    if value == 0.5:
      print("passed.")
    else:
      print("failed.")
    print('P(G|S=s) = {}'.format(value))

    print("Test 3 ....", end = '')    
    factor = restrict_factor(FS, S, '-s')
    factor = restrict_factor(factor, E, '-e')
    factor = restrict_factor(factor, B, 'b')
    value = factor.get_value_at_current_assignments()
    if value == .2:
      print("passed.")
    else:
      print("failed.")
    print('P(S=-s|E=-e,B=b) = {}'.format(value))
    
    print("\nNormalize Tests")
    print("Test 1 ....", end = '')
    normalized_nums = normalize([i for i in range(5)])
    norm_sum = sum(normalized_nums)
    if norm_sum == 1:
      print("passed.")
    else:
      print("failed.")
    print('{} when normalized to {} sum to {}'.format([i for i in range(5)], normalized_nums, norm_sum))

    print("Test 2 ....", end = '')
    normalized_nums = normalize([i for i in range(0,-5,-1)])
    norm_sum = sum(normalized_nums)
    if norm_sum == 1:
      print("passed.")
    else:
      print("failed.")
    print('{} when normalized to {} sum to {}'.format([i for i in range(0,-5,-1)], normalized_nums, norm_sum))

    print("Test 3 ....", end = '')
    normalized_nums = normalize([i for i in range(4,-5,-1)])
    norm_sum = sum(normalized_nums)
    if norm_sum == 0:
      print("passed.")
    else:
      print("failed.")
    print('{} when normalized to {} sum to {}'.format([i for i in range(4,-5,-1)], normalized_nums, norm_sum))

'''