import pymc3 as pm
import numpy as np
from pymc3.backends.smc_text import TextStage
import pytest
from tempfile import mkdtemp
import shutil
import theano.tensor as tt
import theano

from .helpers import SeededTest


@pytest.mark.xfail(condition=(theano.config.floatX == "float32"), reason="Fails on float32")
class TestSMC(SeededTest):

    def setup_class(self):
        super(TestSMC, self).setup_class()
        self.test_folder = mkdtemp(prefix='ATMIP_TEST')

        self.samples = 1500
        self.chains = 200
        n = 4
        mu1 = np.ones(n) * (1. / 2)
        mu2 = - mu1

        stdev = 0.1
        sigma = np.power(stdev, 2) * np.eye(n)
        isigma = np.linalg.inv(sigma)
        dsigma = np.linalg.det(sigma)

        w1 = stdev
        w2 = (1 - stdev)

        def two_gaussians(x):
            log_like1 = - 0.5 * n * tt.log(2 * np.pi) \
                        - 0.5 * tt.log(dsigma) \
                        - 0.5 * (x - mu1).T.dot(isigma).dot(x - mu1)
            log_like2 = - 0.5 * n * tt.log(2 * np.pi) \
                        - 0.5 * tt.log(dsigma) \
                        - 0.5 * (x - mu2).T.dot(isigma).dot(x - mu2)
            return tt.log(w1 * tt.exp(log_like1) + w2 * tt.exp(log_like2))

        with pm.Model() as self.ATMIP_test:
            X = pm.Uniform('X', lower=-2, upper=2., shape=n)
            llk = pm.Potential('muh', two_gaussians(X))

        self.muref = mu1


    @pytest.mark.parametrize(['cores', 'stage'], [[1, 0], [2, 6]])
    def test_sample_n_core(self, cores, stage):
        step_kwargs = {'homepath': self.test_folder, 'stage': stage}
        with self.ATMIP_test:
            mtrace = pm.sample(draws=self.samples,
                               chains=self.chains,
                               cores=cores,
                               step = pm.SMC(),
                               step_kwargs=step_kwargs)

        x = mtrace.get_values('X')
        mu1d = np.abs(x).mean(axis=0)
        np.testing.assert_allclose(self.muref, mu1d, rtol=0., atol=0.03)
        # Scenario IV Ching, J. & Chen, Y. 2007
        #assert np.round(np.log(self.ATMIP_test.marginal_likelihood)) == -12.0

    def test_stage_handler(self):
        stage_number = -1
        stage_handler = TextStage(self.test_folder)

        step = stage_handler.load_atmip_params(stage_number, model=self.ATMIP_test)
        assert step.stage == stage_number

        corrupted_chains = stage_handler.recover_existing_results(stage_number,
                                                                  self.samples / self.chains,
                                                                  self.chains,
                                                                  step,
                                                                  model=self.ATMIP_test)
        assert len(corrupted_chains) == self.chains

        rtrace = stage_handler.load_result_trace(model=self.ATMIP_test)

    def teardown_class(self):
        shutil.rmtree(self.test_folder)
