import unittest
from numpy import *
from stats_toolkit.random import LatinHypercubeRNG
from stats_toolkit.distributions import *
from stats_toolkit.errors import ImproperBoundsError, \
    UndefinedDistributionError

class LatinHypercubeMCTestCase(unittest.TestCase):
    def make_params_array(self, length=1):
        assert isinstance(length, int)
        params = zeros((length,), dtype=[('uncertainty_type', 'i2'), 
            ('input', 'u4'), ('output', 'u4'), 
            ('amount', 'f4'), ('negative', 'b1'), ('sigma', 'f4'), 
            ('minimum', 'f4'), ('maximum', 'f4')])
        params['minimum'] = params['maximum'] = params['sigma'] = NaN
        return params

    def test_known_seed(self):
        params = self.make_params_array(1)
        params['uncertainty_type'] = UniformUncertainty.id
        params['maximum'] = 1
        params['minimum'] = 0
        params['amount'] = 0.5
        params['sigma'] = NaN
        lhc = LatinHypercubeRNG(params, seed=11111, samples=9)
        self.assertTrue(allclose(lhc.iterate(), array([[0.5]])))
        self.assertTrue(allclose(lhc.iterate(), array([[0.7]])))
        self.assertTrue(allclose(lhc.iterate(), array([[0.5]])))
        self.assertTrue(allclose(lhc.iterate(), array([[0.6]])))

    def test_no_uncertainty_intervals(self):
        # Two-dimensional array check
        params = self.make_params_array(2)
        params['uncertainty_type'] = NoUncertainty.id
        params['amount'] = (1,2)
        lhc = LatinHypercubeRNG(params)
        test_values = array(([1, 1, 1, 1, 1, 1, 1, 1, 1, 1], 
            [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]))
        self.assertTrue(allclose(lhc.hypercube, test_values))
        params['uncertainty_type'] = UndefinedUncertainty.id
        lhc = LatinHypercubeRNG(params)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # One-dimensional array check
        params = self.make_params_array(1)
        params['amount'] = 1
        params['uncertainty_type'] = NoUncertainty.id
        lhc = LatinHypercubeRNG(params)
        test_values = array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))

    def test_lognormal_intervals(self):
        # Two-dimensional array check
        params = self.make_params_array(2)
        params['uncertainty_type'] = LognormalUncertainty.id
        params['amount'] = (1,2)
        params['sigma'] = (1,2)
        lhc = LatinHypercubeRNG(params)
        test_values = array([[0.26311141,0.40314545,0.54630091,0.70556548, 
            0.89209264,1.12095981,1.41730289,1.83049303,2.48049433,3.8006714],
            [0.52622281,0.8062909, 1.09260181,1.41113097,1.78418529,2.24191963, 
            2.83460579,3.66098607,4.96098868,7.60134282]])
        # print lhc.hypercube
        # print test_values
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # One-dimensional array check
        params = self.make_params_array(1)
        params['amount'] = 1
        params['sigma'] = 1
        params['uncertainty_type'] = LognormalUncertainty.id
        lhc = LatinHypercubeRNG(params)
        test_values = array([0.26311141,0.40314545,0.54630091,0.70556548, 
            0.89209264,1.12095981,1.41730289,1.83049303,2.48049433,3.8006714] 
            ).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # Lower bound check
        params['amount'] = 1
        params['minimum'] = 1
        lhc = LatinHypercubeRNG(params)
        test_values = array([[1.12095981,1.25845417,1.41730289,1.604463,
            1.83049303,2.11247151,2.48049433,2.99457872,3.8006714,5.42285066] 
            ]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # Upper bound check
        params['amount'] = 1
        params['maximum'] = 1
        params['minimum'] = NaN
        lhc = LatinHypercubeRNG(params)
        test_values = array([[0.18440486,0.26311141,0.33393679,0.40314545,
            0.47337916,0.54630091,0.62326149,0.70556548,0.79462568, 
            0.89209264]]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # Both bounds check
        params['amount'] = 1
        params['minimum'] = 0.5
        params['maximum'] = 1.5
        lhc = LatinHypercubeRNG(params)
        test_values = array([[0.56110473,0.62514369,0.69285366,0.76502388,
            0.84254518,0.92646272,1.01804071,1.11884934,1.23088804,1.35676875] 
            ]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))

    def test_normal_intervals(self):
        # Two-dimensional array check
        params = self.make_params_array(2)
        params['uncertainty_type'] = NormalUncertainty.id
        params['amount'] = (1,2)
        params['sigma'] = (1,2)
        lhc = LatinHypercubeRNG(params)
        test_values = array(([-0.33517774,  0.09154213,  0.39541465,  0.6512443, 
        0.88581471,  1.11418529,  1.3487557 ,  1.60458535,  1.90845787, 
        2.33517774], [-0.67035547,  0.18308426,  0.79082931,  1.30248861, 
        1.77162941,  2.22837059,  2.69751139,  3.20917069,  3.81691574, 
        4.67035547]))
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # One-dimensional array check
        params = self.make_params_array(1)
        params['amount'] = 1
        params['sigma'] = 1
        params['uncertainty_type'] = NormalUncertainty.id
        lhc = LatinHypercubeRNG(params)
        test_values = array([-0.33517774,  0.09154213,  0.39541465, 
            0.6512443, 0.88581471,  1.11418529,  1.3487557 ,  1.60458535, 
            1.90845787, 2.33517774]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # Lower bound check
        params['minimum'] = 1
        lhc = LatinHypercubeRNG(params)
        test_values = array([[1.11418529, 1.22988412, 1.3487557, 1.47278912, 
            1.60458535, 1.74785859, 1.90845787, 2.09680356, 2.33517774, 
            2.69062163]]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # Upper bound check
        params['maximum'] = 1
        params['minimum'] = NaN
        lhc = LatinHypercubeRNG(params)
        test_values = array([[-0.69062163, -0.33517774, -0.09680356, 
            0.09154213, 0.25214141, 0.39541465, 0.52721088, 0.6512443, 
            0.77011588, 0.88581471]]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))
        # Both bounds check
        params['minimum'] = 0.5
        params['maximum'] = 1.5
        lhc = LatinHypercubeRNG(params)
        test_values = array([[0.59665956, 0.68968379, 0.78009253, 0.86873532, 
            0.95635658, 1.04364342, 1.13126468, 1.21990747, 1.31031621, 
            1.40334044]]).reshape(1,10)
        self.assertTrue(allclose(lhc.hypercube, test_values))

    def test_uniform_intervals(self):
        # Two-dimensional array check
        params = self.make_params_array(2)
        params['uncertainty_type'] = UniformUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.5
        params['sigma'] = NaN
        lhc = LatinHypercubeRNG(params, samples=9)
        self.assertTrue(allclose(lhc.hypercube, array([[1.1, 1.2, 1.3, 1.4, 
            1.5, 1.6, 1.7, 1.8, 1.9], [1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 
            1.9]])))
        # One-dimensional array check
        params = self.make_params_array(1)
        params['uncertainty_type'] = UniformUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.5
        params['sigma'] = NaN
        lhc = LatinHypercubeRNG(params, samples=9)
        self.assertTrue(allclose(lhc.hypercube, 
            array([[1.1, 1.2, 1.3, 1.4, 1.5, 
            1.6, 1.7, 1.8, 1.9]])))

    def test_triangular_intervals(self):
        # Two-dimensional array check
        params = self.make_params_array(2)
        params['uncertainty_type'] = TriangularUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.6
        params['sigma'] = NaN
        lhc = LatinHypercubeRNG(params, samples=9)
        self.assertTrue(allclose(lhc.hypercube, array([[1.24494898,1.34641017, 
            1.42426408,1.48989796,1.54772257,1.60000001,1.65358985,1.7171573, 
            1.80000001],[1.24494898,1.34641017,1.42426408,1.48989796, 
            1.54772257, 1.60000001,1.65358985,1.7171573,1.80000001]])))
        # One-dimensional array check
        params = self.make_params_array(1)
        params['uncertainty_type'] = TriangularUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.6
        params['sigma'] = NaN
        lhc = LatinHypercubeRNG(params, samples=9)
        self.assertTrue(allclose(lhc.hypercube, array([[1.24494898,1.34641017, 
            1.42426408,1.48989796,1.54772257,1.60000001,1.65358985,1.7171573, 
            1.80000001]])))

    def test_bernoulli(self):
        # Two-dimensional array check
        params = self.make_params_array(2)
        params['uncertainty_type'] = BernoulliUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.6
        lhc = LatinHypercubeRNG(params)
        self.assertTrue(allclose(lhc.hypercube, array([[1,1,1,1,1,1,2,2,2,2], 
            [1,1,1,1,1,1,2,2,2,2]])))
        # One-dimensional array check
        params = self.make_params_array(1)
        params['uncertainty_type'] = BernoulliUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.6
        lhc = LatinHypercubeRNG(params)
        self.assertTrue(allclose(lhc.hypercube, array([1,1,1,1,1,1,2,2,2,2])))

    def test_heterogeneous_params(self):
        params = self.make_params_array(2)
        params['uncertainty_type'][0] = TriangularUncertainty.id
        params['uncertainty_type'][1] = UniformUncertainty.id
        params['maximum'] = 2
        params['minimum'] = 1
        params['amount'] = 1.6
        params['sigma'] = NaN
        lhc = LatinHypercubeRNG(params, samples=9)
        self.assertTrue(allclose(lhc.hypercube, array([[1.24494898,1.34641017, 
            1.42426408,1.48989796,1.54772257,1.60000001,1.65358985,1.7171573, 
            1.80000001], [1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9]])))
