import pickle
import unittest

import scipy as sp

from util.random_variables import MultivariateMultiTypeNormalDistribution, EmptyMultivariateMultiTypeNormalDistribution, NonEmptyMultivariateMultiTypeNormalDistribution


class TestMultivariateMultiTypeNormalDistribution(unittest.TestCase):
    def test_zero_covariance(self):
        cov = sp.asarray([[0, 0],
                          [0, 0]])
        par_names = ['a', 'b']
        par_types = [float, float]
        zcvs = 2.3
        m = MultivariateMultiTypeNormalDistribution(cov, par_names, par_types,
                                                    zero_covariance_substitutes=zcvs)
        expected = zcvs * sp.eye(2)
        self.assertLess(sp.absolute(expected-m.covariance_matrix).sum(), 1e-5)

    def test_pickleable(self):
        cov = sp.asarray([[0, 0],
                          [0, 0]])

        par_names = ['a', 'b']
        par_types = [float, float]
        zcvs = 2.3
        m = MultivariateMultiTypeNormalDistribution(cov, par_names, par_types,
                                                    zero_covariance_substitutes=zcvs)
        s = pickle.dumps(m)
        m_loaded = pickle.loads(s)
        self.assertEqual(m.parameter_names, m_loaded.parameter_names)
        self.assertEqual(m.parameter_types, m_loaded.parameter_types)
        self.assertTrue((m.covariance_matrix == m_loaded.covariance_matrix).all())
        self.assertIsInstance(m_loaded, NonEmptyMultivariateMultiTypeNormalDistribution)

    def test_empty_multivatiate_pickleable(self):
        cov = sp.asarray([[]])

        par_names = []
        par_types = []
        zcvs = 2.3
        m = MultivariateMultiTypeNormalDistribution(cov, par_names, par_types,
                                                    zero_covariance_substitutes=zcvs)
        s = pickle.dumps(m)
        m_loaded = pickle.loads(s)
        self.assertIsInstance(m_loaded, EmptyMultivariateMultiTypeNormalDistribution)

    def test_multiply(self):
        cov = sp.asarray([[1.1, 0],
                          [0, 2.2]])

        par_names = ['a', 'b']
        par_types = [float, float]
        m = MultivariateMultiTypeNormalDistribution(cov, par_names, par_types)
        m2 = m * 2
        self.assertEqual(1.1, m.covariance_matrix[0,0])
        self.assertEqual(2.2, m.covariance_matrix[1,1])
        self.assertEqual(2.2, m2.covariance_matrix[0,0])
        self.assertEqual(4.4, m2.covariance_matrix[1,1])


if __name__ == '__main__':
    unittest.main()
