from importlib import resources
import pickle
import sys

import numpy as np
import scipy.optimize as sciopt
from scipy.stats import norm

import pytest
from numpy.testing import (
    assert_,
    assert_almost_equal,
    assert_equal,
    assert_allclose,
)

from refnx.analysis import (
    CurveFitter,
    Parameter,
    Parameters,
    Model,
    Objective,
    process_chain,
    load_chain,
    Bounds,
    PDF,
    autocorrelation_chain,
    integrated_time,
)
import refnx.analysis.tests
from refnx.analysis.curvefitter import bounds_list
from refnx.dataset import Data1D
from refnx._lib import emcee, flatten
from refnx._lib.emcee.moves import DEMove, DESnookerMove

from .NISTModels import NIST_runner, NIST_Models


def line(x, params, *args, **kwds):
    p_arr = np.array(params)
    return p_arr[0] + x * p_arr[1]


class TestCurveFitter:
    def setup_method(self):
        # Reproducible results!
        np.random.seed(123)

        self.m_true = -0.9594
        self.b_true = 4.294
        self.f_true = 0.534
        self.m_ls = -1.1040757010910947
        self.b_ls = 5.4405552502319505

        # Generate some synthetic data from the model.
        N = 50
        x = np.sort(10 * np.random.rand(N))
        y_err = 0.1 + 0.5 * np.random.rand(N)
        y = self.m_true * x + self.b_true
        y += np.abs(self.f_true * y) * np.random.randn(N)
        y += y_err * np.random.randn(N)

        self.data = Data1D(data=(x, y, y_err))

        self.p = Parameter(self.b_ls, "b", vary=True, bounds=(-100, 100))
        self.p |= Parameter(self.m_ls, "m", vary=True, bounds=(-100, 100))

        self.model = Model(self.p, fitfunc=line)
        self.objective = Objective(self.model, self.data)
        assert len(self.objective.varying_parameters()) == 2

        mod = np.array(
            [
                4.78166609,
                4.42364699,
                4.16404064,
                3.50343504,
                3.4257084,
                2.93594347,
                2.92035638,
                2.67533842,
                2.28136038,
                2.19772983,
                1.99295496,
                1.93748334,
                1.87484436,
                1.65161016,
                1.44613461,
                1.11128101,
                1.04584535,
                0.86055984,
                0.76913963,
                0.73906649,
                0.73331407,
                0.68350418,
                0.65216599,
                0.59838566,
                0.13070299,
                0.10749131,
                -0.01010195,
                -0.10010155,
                -0.29495372,
                -0.42817431,
                -0.43122391,
                -0.64637715,
                -1.30560686,
                -1.32626428,
                -1.44835768,
                -1.52589881,
                -1.56371158,
                -2.12048349,
                -2.24899179,
                -2.50292682,
                -2.53576659,
                -2.55797996,
                -2.60870542,
                -2.7074727,
                -3.93781479,
                -4.12415366,
                -4.42313742,
                -4.98368609,
                -5.38782395,
                -5.44077086,
            ]
        )
        self.mod = mod

        self.mcfitter = CurveFitter(self.objective)

    def test_bounds_list(self):
        bnds = bounds_list(self.p)
        assert_allclose(bnds, [(-100, 100), (-100, 100)])

        # try making a Parameter bound a normal distribution, then get an
        # approximation to box bounds
        self.p[0].bounds = PDF(norm(0, 1))
        assert_allclose(
            bounds_list(self.p), [norm(0, 1).ppf([0.005, 0.995]), (-100, 100)]
        )

    def test_constraints(self):
        # constraints should work during fitting
        self.p[0].value = 5.4

        self.p[1].constraint = -0.203 * self.p[0]
        assert_equal(self.p[1].value, self.p[0].value * -0.203)
        res = self.mcfitter.fit()

        assert_(res.success)
        assert_equal(len(self.objective.varying_parameters()), 1)

        # lnsigma is parameters[0]
        assert_(self.p[0] is self.objective.parameters.flattened()[0])
        assert_(self.p[1] is self.objective.parameters.flattened()[1])
        assert_almost_equal(self.p[0].value, res.x[0])
        assert_almost_equal(self.p[1].value, self.p[0].value * -0.203)

        # check that constraints work during sampling
        # the CurveFitter has to be set up again if you change how the
        # parameters are being fitted.
        mcfitter = CurveFitter(self.objective)
        assert_(mcfitter.nvary == 1)
        mcfitter.sample(5)
        assert_equal(self.p[1].value, self.p[0].value * -0.203)
        # the constrained parameters should have a chain
        assert_(self.p[0].chain is not None)
        assert_(self.p[1].chain is not None)
        assert_allclose(self.p[1].chain, self.p[0].chain * -0.203)

    def test_mcmc(self):
        self.mcfitter.sample(steps=50, nthin=1, verbose=False)

        assert_equal(self.mcfitter.nvary, 2)

        # smoke test for corner plot
        # for some reason GHA-linux doesn't like this
        self.mcfitter.objective.corner()

        # we're not doing Parallel Tempering here.
        assert_(self.mcfitter._ntemps == -1)
        assert_(isinstance(self.mcfitter.sampler, emcee.EnsembleSampler))

        # should be able to multithread
        mcfitter = CurveFitter(self.objective, nwalkers=50)
        res = mcfitter.sample(steps=33, nthin=2, verbose=False, pool=2)

        # check that the autocorrelation function at least runs
        acfs = mcfitter.acf(nburn=10)
        assert_equal(acfs.shape[-1], mcfitter.nvary)

        # check the standalone autocorrelation calculator
        acfs2 = autocorrelation_chain(mcfitter.chain, nburn=10)
        assert_allclose(acfs, acfs2)

        # check integrated_time
        integrated_time(acfs2, tol=5)

        # check chain shape
        assert_equal(mcfitter.chain.shape, (33, 50, 2))
        # assert_equal(mcfitter._lastpos, mcfitter.chain[:, -1, :])
        assert_equal(res[0].chain.shape, (33, 50))

        # if the number of parameters changes there should be an Exception
        # raised
        from pytest import raises

        with raises(RuntimeError):
            self.p[0].vary = False
            self.mcfitter.sample(1)

        # can fix by making the sampler again
        self.mcfitter.make_sampler()
        self.mcfitter.sample(1)

    def test_moves(self):
        # check that we can provide non-default moves
        m0 = DEMove()
        m1 = DESnookerMove()
        moves = [(m0, 0.8), (m1, 0.2)]
        mcfitter = CurveFitter(self.objective, moves=moves)
        mcfitter.sample(10)

    def test_random_seed(self):
        # check that MCMC sampling is reproducible
        self.mcfitter.sample(steps=2, random_state=1)

        # get a starting pos
        starting_pos = self.mcfitter._state.coords

        # is sampling reproducible
        self.mcfitter.reset()
        self.mcfitter.initialise(pos=starting_pos)
        self.mcfitter.sample(3, random_state=1, pool=1)
        chain1 = np.copy(self.mcfitter.chain)

        self.mcfitter.reset()
        self.mcfitter.initialise(pos=starting_pos)
        self.mcfitter.sample(3, random_state=1, pool=1)
        chain2 = np.copy(self.mcfitter.chain)

        assert_equal(chain1, chain2)

    def test_mcmc_pt(self):
        # smoke test for parallel tempering
        x = np.array(self.objective.parameters)

        mcfitter = CurveFitter(self.objective, ntemps=10, nwalkers=50)
        assert_equal(mcfitter.sampler.ntemps, 10)

        assert len(list(flatten(self.objective.parameters))) == 2
        # check that the parallel sampling works
        # and that chain shape is correct
        res = mcfitter.sample(steps=5, nthin=2, verbose=False, pool=-1)
        assert_equal(mcfitter.chain.shape, (5, 10, 50, 2))
        assert_equal(res[0].chain.shape, (5, 50))
        assert_equal(mcfitter.chain[:, 0, :, 0], res[0].chain)
        assert_equal(mcfitter.chain[:, 0, :, 1], res[1].chain)
        chain = np.copy(mcfitter.chain)

        assert len(list(flatten(self.objective.parameters))) == 2

        # the sampler should store the probability
        assert_equal(mcfitter.logpost.shape, (5, 10, 50))
        assert_allclose(mcfitter.logpost, mcfitter.sampler._ptchain.logP)

        logprobs = mcfitter.logpost
        highest_prob_loc = np.argmax(logprobs[:, 0])
        idx = np.unravel_index(highest_prob_loc, logprobs[:, 0].shape)
        idx = list(idx)
        idx.insert(1, 0)
        idx = tuple(idx)
        assert_equal(idx, mcfitter.index_max_prob)
        pvals = mcfitter.chain[idx]
        assert_allclose(logprobs[idx], self.objective.logpost(pvals))

        # try resetting the chain
        mcfitter.reset()

        # test for reproducible operation
        self.objective.setp(x)
        mcfitter = CurveFitter(self.objective, ntemps=10, nwalkers=50)
        mcfitter.initialise("jitter", random_state=1)
        mcfitter.sample(steps=5, nthin=2, verbose=False, random_state=2)
        chain = np.copy(mcfitter.chain)

        self.objective.setp(x)
        mcfitter = CurveFitter(self.objective, ntemps=10, nwalkers=50)
        mcfitter.initialise("jitter", random_state=1)
        mcfitter.sample(steps=5, nthin=2, verbose=False, random_state=2)
        chain2 = np.copy(mcfitter.chain)

        assert_allclose(chain2, chain)

    def test_mcmc_init(self):
        # smoke test for sampler initialisation
        # TODO check that the initialisation worked.
        # reproducible initialisation with random_state dependents
        self.mcfitter.initialise("prior", random_state=1)
        starting_pos = np.copy(self.mcfitter._state.coords)

        self.mcfitter.initialise("prior", random_state=1)
        starting_pos2 = self.mcfitter._state.coords
        assert_equal(starting_pos, starting_pos2)

        self.mcfitter.initialise("jitter", random_state=1)
        starting_pos = np.copy(self.mcfitter._state.coords)

        self.mcfitter.initialise("jitter", random_state=1)
        starting_pos2 = self.mcfitter._state.coords
        assert_equal(starting_pos, starting_pos2)

        mcfitter = CurveFitter(self.objective, nwalkers=100)
        mcfitter.initialise("covar")
        assert_equal(mcfitter._state.coords.shape, (100, 2))
        mcfitter.initialise("prior")
        assert_equal(mcfitter._state.coords.shape, (100, 2))
        mcfitter.initialise("jitter")
        assert_equal(mcfitter._state.coords.shape, (100, 2))
        # initialise with last position
        mcfitter.sample(steps=1)
        chain = mcfitter.chain
        mcfitter.initialise(pos=chain[-1])
        assert_equal(mcfitter._state.coords.shape, (100, 2))
        # initialise with chain
        mcfitter.sample(steps=2)
        chain = mcfitter.chain
        mcfitter.initialise(pos=chain)
        assert_equal(mcfitter._state.coords, chain[-1])
        # initialise with chain if it's never been run before
        mcfitter = CurveFitter(self.objective, nwalkers=100)
        mcfitter.initialise(chain)

        # initialise for Parallel tempering
        mcfitter = CurveFitter(self.objective, ntemps=20, nwalkers=100)
        mcfitter.initialise("covar")
        assert_equal(mcfitter._state.coords.shape, (20, 100, 2))
        mcfitter.initialise("prior")
        assert_equal(mcfitter._state.coords.shape, (20, 100, 2))
        mcfitter.initialise("jitter")
        assert_equal(mcfitter._state.coords.shape, (20, 100, 2))
        # initialise with last position
        mcfitter.sample(steps=1)
        chain = mcfitter.chain
        mcfitter.initialise(pos=chain[-1])
        assert_equal(mcfitter._state.coords.shape, (20, 100, 2))
        # initialise with chain
        mcfitter.sample(steps=2)
        chain = mcfitter.chain
        mcfitter.initialise(pos=np.copy(chain))
        assert_equal(mcfitter._state.coords, chain[-1])
        # initialise with chain if it's never been run before
        mcfitter = CurveFitter(self.objective, nwalkers=100, ntemps=20)
        mcfitter.initialise(chain)

    def test_fit_smoke(self):
        # smoke tests to check that fit runs
        def callback(xk):
            return

        def callback2(xk, *args, **kws):
            return

        # L-BFGS-B
        res0 = self.mcfitter.fit(callback=callback)
        assert_almost_equal(res0.x, [self.b_ls, self.m_ls], 6)
        res0 = self.mcfitter.fit()
        res0 = self.mcfitter.fit(verbose=False)
        res0 = self.mcfitter.fit(verbose=False, callback=callback)

        # least_squares
        res1 = self.mcfitter.fit(method="least_squares")
        assert_almost_equal(res1.x, [self.b_ls, self.m_ls], 6)

        # least_squares doesn't accept a callback. As well as testing that
        # least_squares works, it checks that providing a callback doesn't
        # trip the fitter up.
        res1 = self.mcfitter.fit(method="least_squares", callback=callback)
        assert_almost_equal(res1.x, [self.b_ls, self.m_ls], 6)

        # need full bounds for differential_evolution
        self.p[0].range(3, 7)
        self.p[1].range(-2, 0)
        res2 = self.mcfitter.fit(
            method="differential_evolution",
            seed=1,
            popsize=10,
            maxiter=100,
            callback=callback2,
        )
        assert_almost_equal(res2.x, [self.b_ls, self.m_ls], 6)

        # check that the res object has covar and stderr
        assert_("covar" in res0)
        assert_("stderr" in res0)

    @pytest.mark.parametrize("model", NIST_Models)
    def test_NIST(self, model):
        # Run all the NIST standard tests with leastsq
        NIST_runner(model)


def gauss(x, p0):
    p = np.array(p0)
    return p[0] + p[1] * np.exp(-(((x - p[2]) / p[3]) ** 2))


class TestFitterGauss:
    # Test CurveFitter with a noisy gaussian, weighted and unweighted, to see
    # if the parameters and uncertainties come out correct

    @pytest.fixture(autouse=True)
    def setup_method(self, tmp_path):
        self.path = resources.files(refnx.analysis.tests)
        self.tmp_path = tmp_path

        theoretical = np.loadtxt(self.path / "gauss_data.txt")
        xvals, yvals, evals = np.hsplit(theoretical, 3)
        xvals = xvals.flatten()
        yvals = yvals.flatten()
        evals = evals.flatten()

        # these best weighted values and uncertainties obtained with Igor
        self.best_weighted = [-0.00246095, 19.5299, -8.28446e-2, 1.24692]

        self.best_weighted_errors = [
            0.0220313708486,
            1.12879436221,
            0.0447659158681,
            0.0412022938883,
        ]

        self.best_weighted_chisqr = 77.6040960351

        self.best_unweighted = [
            -0.10584111872702096,
            19.240347049328989,
            0.0092623066070940396,
            1.501362314145845,
        ]

        self.best_unweighted_errors = [
            0.34246565477,
            0.689820935208,
            0.0411243173041,
            0.0693429375282,
        ]

        self.best_unweighted_chisqr = 497.102084956

        self.p0 = np.array([0.1, 20.0, 0.1, 0.1])
        self.names = ["bkg", "A", "x0", "width"]
        self.bounds = [(-1, 1), (0, 30), (-5.0, 5.0), (0.001, 2)]

        self.params = Parameters(name="gauss_params")
        for p, name, bound in zip(self.p0, self.names, self.bounds):
            param = Parameter(p, name=name)
            param.range(*bound)
            param.vary = True
            self.params.append(param)

        self.model = Model(self.params, fitfunc=gauss)
        self.data = Data1D((xvals, yvals, evals))
        self.objective = Objective(self.model, self.data)
        return 0

    def test_pickle(self):
        # tests if a CurveFitter can be pickled/unpickled.
        f = CurveFitter(self.objective)
        pkl = pickle.dumps(f)
        g = pickle.loads(pkl)
        g._check_vars_unchanged()

    def test_best_weighted(self):
        assert_equal(len(self.objective.varying_parameters()), 4)
        self.objective.setp(self.p0)

        f = CurveFitter(self.objective, nwalkers=100)
        res = f.fit("least_squares", jac="3-point")

        output = res.x
        assert_almost_equal(output, self.best_weighted, 3)
        assert_almost_equal(
            self.objective.chisqr(), self.best_weighted_chisqr, 5
        )

        # compare the residuals
        res = (self.data.y - self.model(self.data.x)) / self.data.y_err
        assert_equal(self.objective.residuals(), res)

        # compare objective.covar to the best_weighted_errors
        uncertainties = [param.stderr for param in self.params]
        assert_allclose(uncertainties, self.best_weighted_errors, rtol=0.005)

        # we're also going to try the checkpointing here.
        checkpoint = self.tmp_path / "checkpoint.txt"

        # compare samples to best_weighted_errors
        np.random.seed(1)
        f.sample(steps=201, random_state=1, verbose=False, f=checkpoint)
        process_chain(self.objective, f.chain, nburn=50, nthin=10)
        uncertainties = [param.stderr for param in self.params]
        assert_allclose(uncertainties, self.best_weighted_errors, rtol=0.07)

        # test that the checkpoint worked
        check_array = np.loadtxt(checkpoint)
        check_array = check_array.reshape(201, f._nwalkers, f.nvary)
        assert_allclose(check_array, f.chain)

        # test loading the checkpoint
        chain = load_chain(checkpoint)
        assert_allclose(chain, f.chain)

        f.initialise("jitter")
        f.sample(steps=2, nthin=4, f=checkpoint, verbose=False)
        assert_equal(f.chain.shape[0], 2)

        # we should be able to produce 2 * 100 steps from the generator
        g = self.objective.pgen(ngen=20000000000)
        s = [i for i, a in enumerate(g)]
        assert_equal(np.max(s), 200 - 1)
        g = self.objective.pgen(ngen=200)
        pvec = next(g)
        assert_equal(pvec.size, len(self.objective.parameters.flattened()))

        # check that all the parameters are returned via pgen, not only those
        # being varied.
        self.params[0].vary = False
        f = CurveFitter(self.objective, nwalkers=100)
        f.initialise("jitter")
        f.sample(steps=2, nthin=4, f=checkpoint, verbose=False)
        g = self.objective.pgen(ngen=100)
        pvec = next(g)
        assert_equal(pvec.size, len(self.objective.parameters.flattened()))

        # the following test won't work because of emcee/gh226.
        # chain = load_chain(checkpoint)
        # assert_(chain.shape == f.chain.shape)
        # assert_allclose(chain, f.chain)

        # try reproducing best fit with parallel tempering
        self.params[0].vary = True
        f = CurveFitter(self.objective, nwalkers=100, ntemps=10)
        f.fit("differential_evolution", seed=1)

        f.sample(steps=801, random_state=2, verbose=False)
        assert isinstance(
            f.sampler._ptchain.ensemble._rng, np.random.Generator
        )

        process_chain(self.objective, f.chain, nburn=100, nthin=20)
        print(self.params[0].chain.shape, self.params[0].chain)

        uncertainties = [param.stderr for param in self.params]
        assert_allclose(uncertainties, self.best_weighted_errors, rtol=0.09)

    def test_best_unweighted(self):
        self.objective.weighted = False
        f = CurveFitter(self.objective, nwalkers=100)
        res = f.fit()

        output = res.x
        assert_almost_equal(
            self.objective.chisqr(), self.best_unweighted_chisqr
        )
        assert_almost_equal(output, self.best_unweighted, 5)

        # compare the residuals
        res = self.data.y - self.model(self.data.x)
        assert_equal(self.objective.residuals(), res)

        # compare objective._covar to the best_unweighted_errors
        uncertainties = np.array([param.stderr for param in self.params])
        assert_almost_equal(uncertainties, self.best_unweighted_errors, 3)

        # the samples won't compare to the covariance matrix...
        # f.sample(nsteps=150, nburn=20, nthin=30, random_state=1)
        # uncertainties = [param.stderr for param in self.params]
        # assert_allclose(uncertainties, self.best_unweighted_errors,
        #                 rtol=0.15)

    def test_all_minimisers(self):
        """test minimisers against the Gaussian fit"""
        f = CurveFitter(self.objective)

        methods = [
            "differential_evolution",
            "L-BFGS-B",
            "least_squares",
            "shgo",
            "dual_annealing",
        ]

        for method in methods:
            self.objective.setp(self.p0)
            opts = {}
            if method in ["differential_evolution", "dual_annealing"]:
                opts = {"seed": 1}

            res = f.fit(method=method, **opts)
            assert_allclose(res.x, self.best_weighted, rtol=0.005)

        # smoke test to check that we can use nlpost
        self.objective.setp(self.p0)
        logp0 = self.objective.logp()

        # check that probabilities are calculated correctly
        assert_allclose(
            self.objective.logpost(),
            self.objective.logp() + self.objective.logl(),
        )
        assert_allclose(self.objective.nlpost(), -self.objective.logpost())
        assert_allclose(
            self.objective.nlpost(self.p0), -self.objective.logpost(self.p0)
        )

        # if the priors are all uniform then the only difference between
        # logpost and logl is a constant. A minimiser should converge on the
        # same answer. The following tests examine that.
        # The test works for dual_annealing, but not for differential
        # evolution, not sure why that is.
        self.objective.setp(self.p0)
        res1 = f.fit(method="dual_annealing", seed=1)
        assert_almost_equal(res1.x, self.best_weighted, 3)
        nll1 = self.objective.nll()
        nlpost1 = self.objective.nlpost()

        self.objective.setp(self.p0)
        res2 = f.fit(method="dual_annealing", target="nlpost", seed=1)
        assert_almost_equal(res2.x, self.best_weighted, 3)
        nll2 = self.objective.nll()
        nlpost2 = self.objective.nlpost()

        assert_allclose(nlpost1, nlpost2, atol=0.001)
        assert_allclose(nll1, nll2, atol=0.001)

        # these two priors are calculated for different parameter values
        # (before and after the fit) they should be the same because all
        # the parameters have uniform priors.
        assert_almost_equal(self.objective.logp(), logp0)

    @pytest.mark.xfail(
        sys.platform == "win32"
        and (sys.version_info.major, sys.version_info.minor)
        in [(3, 8), (3, 9), (3, 10)],
        run=False,
        reason="doesn't work on cp38",
    )
    def test_pymc_sample(self):
        # test sampling with pymc
        try:
            import pymc as pm
            from refnx.analysis import pymc_model
        except (ModuleNotFoundError, ImportError, AttributeError):
            # can't run test if pymc/pytensor not installed
            return

        with pymc_model(self.objective):
            s = pm.NUTS()
            pm.sample(
                200,
                tune=100,
                step=s,
                discard_tuned_samples=True,
                compute_convergence_checks=False,
                random_seed=1,
            )
            # don't check for accuracy because it requires a heap more
            # draws.
            # means = [np.mean(trace[f"p{i}"]) for i in range(4)]
            # assert_allclose(means, self.best_weighted, rtol=0.04)
            # errors = [np.std(trace[f"p{i}"]) for i in range(4)]
            # assert_allclose(errors, self.best_weighted_errors, atol=0.02)


"""
        The Gaussian example sampling can also be performed with pymc.
        The above results from emcee have been verified against pymc - the
        unweighted sampling statistics are the same.

        from pymc import (Model, Normal, HalfNormal, Flat, Uniform,
                          find_MAP, NUTS, sample, summary, traceplot)

        basic_model = Model()

        with basic_model:
            # Priors for unknown model parameters
            bkg = Uniform('bkg', -1, 5)
            A0 = Uniform('A0', 0, 50)
            x0 = Uniform('x0', min(x), max(x))
            width = Uniform('width', 0.5, 10)

            # Expected value of outcome
            mu = bkg + A0 * np.exp(-((x - x0) / width) ** 2)

            # Likelihood (sampling distribution) of observations
            #     y_obs = Normal('y_obs', mu=mu, sd=e, observed=y)
            y_obs = Normal('y_obs', mu=mu, observed=y)

        with basic_model:
            # draw 500 posterior samples
            trace = sample(500)
        summary(trace)
"""
