import numpy as np
try:
    import unittest.mock as mock  # py3
except ImportError:
    import mock
import unittest

import pymc3 as pm
import theano.tensor as tt
from theano import shared
from .models import simple_init
from .helpers import SeededTest

# Test if multiprocessing is available
import multiprocessing
try:
    multiprocessing.Pool(2)
except:
    pass


class TestSample(SeededTest):
    def setUp(self):
        super(TestSample, self).setUp()
        self.model, self.start, self.step, _ = simple_init()

    def test_sample_does_not_set_seed(self):
        random_numbers = []
        for _ in range(2):
            np.random.seed(1)
            with self.model:
                pm.sample(1)
                random_numbers.append(np.random.random())
        self.assertEqual(random_numbers[0], random_numbers[1])

    def test_sample(self):
        test_njobs = [1]
        with self.model:
            for njobs in test_njobs:
                for steps in [1, 10, 300]:
                    pm.sample(steps, self.step, {}, None, njobs=njobs, random_seed=self.random_seed)

    def test_iter_sample(self):
        with self.model:
            samps = pm.sampling.iter_sample(5, self.step, self.start, random_seed=self.random_seed)
            for i, trace in enumerate(samps):
                self.assertEqual(i, len(trace) - 1, "Trace does not have correct length.")

    def test_parallel_start(self):
        with self.model:
            tr = pm.sample(5, njobs=2, start=[{'x': [10, 10]}, {'x': [-10, -10]}],
                           random_seed=self.random_seed)
        self.assertGreater(tr.get_values('x', chains=0)[0][0], 0)
        self.assertLess(tr.get_values('x', chains=1)[0][0], 0)


class SoftUpdate(SeededTest):
    def test_soft_update_all_present(self):
        start = {'a': 1, 'b': 2}
        test_point = {'a': 3, 'b': 4}
        pm.sampling._soft_update(start, test_point)
        self.assertDictEqual(start, {'a': 1, 'b': 2})

    def test_soft_update_one_missing(self):
        start = {'a': 1, }
        test_point = {'a': 3, 'b': 4}
        pm.sampling._soft_update(start, test_point)
        self.assertDictEqual(start, {'a': 1, 'b': 4})

    def test_soft_update_empty(self):
        start = {}
        test_point = {'a': 3, 'b': 4}
        pm.sampling._soft_update(start, test_point)
        self.assertDictEqual(start, test_point)


class TestNamedSampling(SeededTest):
    def test_shared_named(self):
        G_var = shared(value=np.atleast_2d(1.), broadcastable=(True, False),
                       name="G")

        with pm.Model():
            theta0 = pm.Normal('theta0', mu=np.atleast_2d(0),
                               tau=np.atleast_2d(1e20), shape=(1, 1),
                               testval=np.atleast_2d(0))
            theta = pm.Normal('theta', mu=tt.dot(G_var, theta0),
                              tau=np.atleast_2d(1e20), shape=(1, 1))
            res = theta.random()
            assert np.isclose(res, 0.)

    def test_shared_unnamed(self):
        G_var = shared(value=np.atleast_2d(1.), broadcastable=(True, False))
        with pm.Model():
            theta0 = pm.Normal('theta0', mu=np.atleast_2d(0),
                               tau=np.atleast_2d(1e20), shape=(1, 1),
                               testval=np.atleast_2d(0))
            theta = pm.Normal('theta', mu=tt.dot(G_var, theta0),
                              tau=np.atleast_2d(1e20), shape=(1, 1))
            res = theta.random()
            assert np.isclose(res, 0.)

    def test_constant_named(self):
        G_var = tt.constant(np.atleast_2d(1.), name="G")
        with pm.Model():
            theta0 = pm.Normal('theta0', mu=np.atleast_2d(0),
                               tau=np.atleast_2d(1e20), shape=(1, 1),
                               testval=np.atleast_2d(0))
            theta = pm.Normal('theta', mu=tt.dot(G_var, theta0),
                              tau=np.atleast_2d(1e20), shape=(1, 1))

            res = theta.random()
            assert np.isclose(res, 0.)



class TestChooseBackend(unittest.TestCase):
    def test_choose_backend_none(self):
        with mock.patch('pymc3.sampling.NDArray') as nd:
            pm.sampling._choose_backend(None, 'chain')
        self.assertTrue(nd.called)

    def test_choose_backend_list_of_variables(self):
        with mock.patch('pymc3.sampling.NDArray') as nd:
            pm.sampling._choose_backend(['var1', 'var2'], 'chain')
        nd.assert_called_with(vars=['var1', 'var2'])

    def test_choose_backend_invalid(self):
        self.assertRaises(ValueError,
                          pm.sampling._choose_backend,
                          'invalid', 'chain')

    def test_choose_backend_shortcut(self):
        backend = mock.Mock()
        shortcuts = {'test_backend': {'backend': backend,
                                      'name': None}}
        pm.sampling._choose_backend('test_backend', 'chain', shortcuts=shortcuts)
        self.assertTrue(backend.called)
