#!/usr/bin/env python3
import unittest
import numpy as np
from pgmpy.models import BayesianModel
from pgmpy.models import MarkovModel
from pgmpy.factors import Factor
from pgmpy.factors import TabularCPD
from pgmpy.inference import Inference
from collections import defaultdict


class TestInferenceBase(unittest.TestCase):
    def setUp(self):
        self.bayesian = BayesianModel([('a', 'b'), ('b', 'c'), ('c', 'd'), ('d', 'e')])
        a_cpd = TabularCPD('a', 2, [[0.4, 0.6]])
        b_cpd = TabularCPD('b', 2, [[0.2, 0.4], [0.3, 0.4]], evidence='a',
                           evidence_card=[2])
        c_cpd = TabularCPD('c', 2, [[0.1, 0.2], [0.3, 0.4]], evidence='b',
                           evidence_card=[2])
        d_cpd = TabularCPD('d', 2, [[0.4, 0.3], [0.2, 0.1]], evidence='c',
                           evidence_card=[2])
        e_cpd = TabularCPD('e', 2, [[0.3, 0.2], [0.4, 0.1]], evidence='d',
                           evidence_card=[2])
        self.bayesian.add_cpds(a_cpd, b_cpd, c_cpd, d_cpd, e_cpd)

        self.markov = MarkovModel([('a', 'b'), ('b', 'd'), ('a', 'c'), ('c', 'd')])
        factor_1 = Factor(['a', 'b'], [2, 2], np.array([100, 1, 1, 100]))
        factor_2 = Factor(['a', 'c'], [2, 2], np.array([40, 30, 100, 20]))
        factor_3 = Factor(['b', 'd'], [2, 2], np.array([1, 100, 100, 1]))
        factor_4 = Factor(['c', 'd'], [2, 2], np.array([60, 60, 40, 40]))
        self.markov.add_factors(factor_1, factor_2, factor_3, factor_4)

    def test_bayesian_inference_init(self):
        infer_bayesian = Inference(self.bayesian)
        self.assertEqual(set(infer_bayesian.variables), {'a', 'b', 'c', 'd', 'e'})
        self.assertEqual(infer_bayesian.cardinality, {'a': 2, 'b': 2, 'c': 2,
                                                      'd': 2, 'e': 2})
        self.assertIsInstance(infer_bayesian.factors, defaultdict)
        self.assertEqual(set(infer_bayesian.factors['a']),
                         {[self.bayesian.get_cpds('a').to_factor(),
                           self.bayesian.get_cpds('b').to_factor()]})
        self.assertEqual(set(infer_bayesian.factors['b']),
                         {[self.bayesian.get_cpds('b').to_factor(),
                           self.bayesian.get_cpds('c').to_factor()]})
        self.assertEqual(set(infer_bayesian.factors['c']),
                         {[self.bayesian.get_cpds('c').to_factor(),
                           self.bayesian.get_cpds('d').to_factor()]})
        self.assertEqual(set(infer_bayesian.factors['d']),
                         {[self.bayesian.get_cpds('d').to_factor(),
                           self.bayesian.get_cpds('e').to_factor()]})
        self.assertEqual(set(infer_bayesian.factors['e']),
                         {[self.bayesian.get_cpds('e').to_factor()]})

    def test_markov_inference_init(self):
        infer_markov = Inference(self.markov)
        self.assertEqual(set(infer_markov.variables), {'a', 'b', 'c', 'd'})
        self.assertEqual(infer_markov.cardinality, {'a': 2, 'b': 2, 'c': 2, 'd': 2})
        self.assertEqual(infer_markov.factors, {'a': [Factor(['a', 'b'], [2, 2],
                                                             np.array([100, 1, 1, 100])),
                                                      Factor(['a', 'c'], [2, 2],
                                                             np.array([40, 30, 100, 20]))],
                                                'b': [Factor(['a', 'b'], [2, 2],
                                                             np.array([100, 1, 1, 100])),
                                                      Factor(['b', 'd'], [2, 2],
                                                             np.array([1, 100, 100, 1]))],
                                                'c': [Factor(['a', 'c'], [2, 2],
                                                             np.array([40, 30, 100, 20])),
                                                      Factor(['c', 'd'], [2, 2],
                                                             np.array([60, 60, 40, 40]))],
                                                'd': [Factor(['b', 'd'], [2, 2],
                                                             np.array([1, 100, 100, 1])),
                                                      Factor(['c', 'd'], [2, 2],
                                                             np.array([60, 60, 40, 40]))]})

@@ -1949,33 +1949,36 @@
                 
-%7B
+set(
 %5Bself.bayesian.g
@@ -2022,32 +2022,35 @@
                 
+ 
+  
 self.bayesian.ge
@@ -2065,33 +2065,33 @@
 b').to_factor()%5D
-%7D
+)
 )%0A        self.a
@@ -2152,33 +2152,36 @@
                 
-%7B
+set(
 %5Bself.bayesian.g
@@ -2198,32 +2198,35 @@
 ').to_factor(),%0A
+   
                 
@@ -2268,33 +2268,33 @@
 c').to_factor()%5D
-%7D
+)
 )%0A        self.a
@@ -2355,33 +2355,36 @@
                 
-%7B
+set(
 %5Bself.bayesian.g
@@ -2428,32 +2428,35 @@
                 
+   
 self.bayesian.ge
@@ -2471,33 +2471,33 @@
 d').to_factor()%5D
-%7D
+)
 )%0A        self.a
@@ -2558,33 +2558,36 @@
                 
-%7B
+set(
 %5Bself.bayesian.g
@@ -2631,32 +2631,35 @@
                 
+ 
+  
 self.bayesian.ge
@@ -2682,17 +2682,17 @@
 actor()%5D
-%7D
+)
 )%0A      
@@ -2769,17 +2769,20 @@
         
-%7B
+set(
 %5Bself.ba
@@ -2814,17 +2814,17 @@
 actor()%5D
-%7D
+)
 )%0A%0A    d

Revert "very minor changes"