import unittest

import numpy as np
import numpy.testing as np_test

from pgmpy.factors.distributions import GaussianDistribution as GD


class TestGDInit(unittest.TestCase):
    def test_class_init(self):
        phi1 = GD(['x1', 'x2', 'x3'], np.array([[1], [-3], [4]]),
                  np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertEqual(phi1.variables, ['x1', 'x2', 'x3'])
        np_test.assert_array_equal(phi1.mean, np.asarray([[1], [-3], [4]], dtype=float))
        np_test.assert_array_equal(phi1.covariance,
                                   np.asarray([[4, 2, -2], [2, 5, -5], [-2, -5, 8]], dtype=float))
        self.assertEqual(phi1._precision_matrix, None)

        phi2 = GD(['x1', 'x2', 'x3'], [1, 2, 5],
                  np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertEqual(phi2.variables, ['x1', 'x2', 'x3'])
        np_test.assert_array_equal(phi2.mean, np.asarray([[1], [2], [5]], dtype=float))
        np_test.assert_array_equal(phi2.covariance,
                                   np.asarray([[4, 2, -2], [2, 5, -5], [-2, -5, 8]], dtype=float))
        self.assertEqual(phi2._precision_matrix, None)

        phi3 = GD(['x'], [0], [[1]])
        self.assertEqual(phi3.variables, ['x'])
        np_test.assert_array_equal(phi3.mean, np.asarray([[0]], dtype=float))
        np_test.assert_array_equal(phi3.covariance, np.asarray([[1]], dtype=float))
        self.assertEqual(phi3._precision_matrix, None)

        phi1 = GD(['1', 2, (1, 2, 'x')], np.array([[1], [-3], [4]]),
                   np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertEqual(phi1.variables, ['1', 2, (1, 2, 'x')])
        np_test.assert_array_equal(phi1.mean, np.asarray([[1], [-3], [4]], dtype=float))
        np_test.assert_array_equal(phi1.covariance,
                                   np.asarray([[4, 2, -2], [2, 5, -5], [-2, -5, 8]], dtype=float))
        self.assertEqual(phi1._precision_matrix, None)

        phi2 = GD(['1', 7, (1, 2, 'x')], [1, 2, 5],
                   np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertEqual(phi2.variables, ['1', 7, (1, 2, 'x')])
        np_test.assert_array_equal(phi2.mean, np.asarray([[1], [2], [5]], dtype=float))
        np_test.assert_array_equal(phi2.covariance,
                                   np.asarray([[4, 2, -2], [2, 5, -5], [-2, -5, 8]], dtype=float))
        self.assertEqual(phi2._precision_matrix, None)

        phi3 = GD([23], [0], [[1]])
        self.assertEqual(phi3.variables, [23])
        np_test.assert_array_equal(phi3.mean, np.asarray([[0]], dtype=float))
        np_test.assert_array_equal(phi3.covariance, np.asarray([[1]], dtype=float))
        self.assertEqual(phi3._precision_matrix, None)

    def test_class_init_valueerror(self):
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1, -3],
                          np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2'], [1, -3, 4],
                          np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [[1, -3, 4]],
                          np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1],
                          np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [[1], [-3]],
                          np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))

        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1, -3, 4],
                          np.array([[4, 2, -2], [2, 5, -5]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1, -3, 4],
                          np.array([[4, 2, -2]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1, -3],
                          np.array([[4, 2], [2, 5], [-2, -5]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1, -3],
                          np.array([[4], [2], [-2]]))
        self.assertRaises(ValueError, GD, ['x1', 'x2', 'x3'], [1, -3],
                          np.array([[-2]]))


class TestJGDMethods(unittest.TestCase):
    def setUp(self):
        self.phi1 = GD(['x1', 'x2', 'x3'], np.array([[1], [-3], [4]]),
                        np.array([[4, 2, -2], [2, 5, -5], [-2, -5, 8]]))
        self.phi2 = GD(['x'], [0], [[1]])
        self.phi3 = self.phi1.copy()

    def test_precision_matrix(self):
        self.assertEqual(self.phi1._precision_matrix, None)
        np_test.assert_almost_equal(self.phi1.precision_matrix,
                                    np.array([[0.3125, -0.125, 0],
                                             [-0.125, 0.5833333, 0.3333333],
                                             [0, 0.3333333, 0.3333333]]))
        np_test.assert_almost_equal(self.phi1._precision_matrix,
                                    np.array([[0.3125, -0.125, 0],
                                             [-0.125, 0.5833333, 0.3333333],
                                             [0, 0.3333333, 0.3333333]]))

        self.assertEqual(self.phi2._precision_matrix, None)
        np_test.assert_almost_equal(self.phi2.precision_matrix, np.array([[1]]))
        np_test.assert_almost_equal(self.phi2._precision_matrix, np.array([[1]]))

    def test_marginalize(self):
        phi = self.phi1.marginalize(['x3'], inplace=False)
        self.assertEqual(phi.variables, ['x1', 'x2'])
        np_test.assert_array_equal(phi.mean, np.asarray([[1], [-3]], dtype=float))
        np_test.assert_array_equal(phi.covariance,
                                   np.asarray([[4, 2], [2, 5]], dtype=float))
        self.assertEqual(phi._precision_matrix, None)

        phi = self.phi1.marginalize(['x3', 'x2'], inplace=False)
        self.assertEqual(phi.variables, ['x1'])
        np_test.assert_array_equal(phi.mean, np.asarray([[1]], dtype=float))
        np_test.assert_array_equal(phi.covariance,
                                   np.asarray([[4]], dtype=float))
        self.assertEqual(phi._precision_matrix, None)

        self.phi1.marginalize(['x3'])
        self.assertEqual(self.phi1.variables, ['x1', 'x2'])
        np_test.assert_array_equal(self.phi1.mean, np.asarray([[1], [-3]], dtype=float))
        np_test.assert_array_equal(self.phi1.covariance,
                                   np.asarray([[4, 2], [2, 5]], dtype=float))
        self.assertEqual(self.phi1._precision_matrix, None)

        self.phi1 = self.phi3
        self.phi1.marginalize(['x3', 'x2'])
        self.assertEqual(self.phi1.variables, ['x1'])
        np_test.assert_array_equal(self.phi1.mean, np.asarray([[1]], dtype=float))
        np_test.assert_array_equal(self.phi1.covariance,
                                   np.asarray([[4]], dtype=float))
        self.assertEqual(self.phi1._precision_matrix, None)

        self.phi1 = self.phi3

    def test_copy(self):
        copy_phi1 = self.phi1.copy()
        self.assertEqual(copy_phi1.variables, self.phi1.variables)
        np_test.assert_array_equal(copy_phi1.mean, self.phi1.mean)
        np_test.assert_array_equal(copy_phi1.covariance, self.phi1.covariance)
        np_test.assert_array_equal(copy_phi1._precision_matrix, self.phi1._precision_matrix)

        copy_phi1.marginalize(['x3'])
        self.assertEqual(self.phi1.variables, ['x1', 'x2', 'x3'])
        np_test.assert_array_equal(self.phi1.mean, np.asarray([[1], [-3], [4]], dtype=float))
        np_test.assert_array_equal(self.phi1.covariance,
                                   np.asarray([[4, 2, -2], [2, 5, -5], [-2, -5, 8]], dtype=float))
        self.assertEqual(self.phi1._precision_matrix, None)

        self.phi1.marginalize(['x2'])
        self.assertEqual(copy_phi1.variables, ['x1', 'x2'])
        np_test.assert_array_equal(copy_phi1.mean, np.asarray([[1], [-3]], dtype=float))
        np_test.assert_array_equal(copy_phi1.covariance,
                                   np.asarray([[4, 2], [2, 5]], dtype=float))
        self.assertEqual(copy_phi1._precision_matrix, None)

        self.phi1 = self.phi3

    def test_assignment(self):
        np_test.assert_almost_equal(self.phi1.assignment(*[1, 2, 3]), 2.797826e-05)
        np_test.assert_almost_equal(self.phi1.assignment(*[[1, 2, 3], [0, 0, 0]]),
                                    np.array([2.79782602e-05, 1.48056313e-03]))
        np_test.assert_almost_equal(self.phi2.assignment(0), 0.3989422804)
        np_test.assert_almost_equal(self.phi2.assignment(*[0, 1, -1]),
                                    np.array([0.39894228, 0.24197072, 0.24197072]))

    def test_reduce(self):
        phi = self.phi1.reduce([('x1', 7)], inplace=False)
        self.assertEqual(phi.variables, ['x2', 'x3'])
        np_test.assert_array_equal(phi.mean, np.asarray([[0], [1]], dtype=float))
        np_test.assert_array_equal(phi.covariance,
                                   np.asarray([[4, -4], [-4, 7]], dtype=float))
        self.assertEqual(phi._precision_matrix, None)

        phi = self.phi1.reduce([('x1', 3), ('x2', 1)], inplace=False)
        self.assertEqual(phi.variables, ['x3'])
        np_test.assert_array_equal(phi.mean, np.array([[0]], dtype=float))
        np_test.assert_array_equal(phi.covariance,
                                   np.asarray([[3]], dtype=float))
        self.assertEqual(phi._precision_matrix, None)

        self.phi1.reduce([('x1', 7)])
        self.assertEqual(self.phi1.variables, ['x2', 'x3'])
        np_test.assert_array_equal(self.phi1.mean, np.asarray([[0], [1]], dtype=float))
        np_test.assert_array_equal(self.phi1.covariance,
                                   np.asarray([[4, -4], [-4, 7]], dtype=float))
        self.assertEqual(self.phi1._precision_matrix, None)

        self.phi1 = self.phi3.copy()
        self.phi1.reduce([('x1', 3), ('x2', 1)])
        self.assertEqual(self.phi1.variables, ['x3'])
        np_test.assert_array_equal(self.phi1.mean, np.asarray([[0]], dtype=float))
        np_test.assert_array_equal(self.phi1.covariance,
                                   np.asarray([[3]], dtype=float))
        self.assertEqual(self.phi1._precision_matrix, None)

        self.phi1 = self.phi3.copy()
        self.phi1.reduce([('x2', 1), ('x1', 3)])
        self.assertEqual(self.phi1.variables, ['x3'])
        np_test.assert_array_equal(self.phi1.mean, np.asarray([[0]], dtype=float))
        np_test.assert_array_equal(self.phi1.covariance,
                                   np.asarray([[3]], dtype=float))
        self.assertEqual(self.phi1._precision_matrix, None)

    def test_normalize(self):
        phi = self.phi1.copy()
        phi.normalize()
        self.assertEqual(self.phi1.variables, phi.variables)
        np_test.assert_array_equal(self.phi1.mean, phi.mean)
        np_test.assert_array_equal(self.phi1.covariance, phi.covariance)
        self.assertEqual(self.phi1._precision_matrix, phi._precision_matrix)

        phi = self.phi1.normalize(inplace=False)
        self.assertEqual(self.phi1.variables, phi.variables)
        np_test.assert_array_equal(self.phi1.mean, phi.mean)
        np_test.assert_array_equal(self.phi1.covariance, phi.covariance)
        self.assertEqual(self.phi1._precision_matrix, phi._precision_matrix)

    def test_product(self):
        pass

    def test_divide(self):
        pass

    def test_eq(self):
        pass

    def test_repr(self):
        pass

    def tearDown(self):
        del self.phi1
        del self.phi2
        del self.phi3
