import numpy as np
from numpy.testing import assert_allclose, assert_array_less
import pytest

from scipy import stats
from scipy.stats import sobol_indices
from scipy.stats._resampling import BootstrapResult
from scipy.stats._sensitivity_analysis import (
    BootstrapSobolResult, f_ishigami, sample_AB, sample_A_B
)


@pytest.fixture(scope='session')
def ishigami_ref_indices():
    """Reference values for Ishigami from Saltelli2007.

    Chapter 4, exercise 5 pages 179-182.
    """
    a = 7.
    b = 0.1

    var = 0.5 + a**2/8 + b*np.pi**4/5 + b**2*np.pi**8/18
    v1 = 0.5 + b*np.pi**4/5 + b**2*np.pi**8/50
    v2 = a**2/8
    v3 = 0
    v12 = 0
    # v13: mistake in the book, see other derivations e.g. in 10.1002/nme.4856
    v13 = b**2*np.pi**8*8/225
    v23 = 0

    s_first = np.array([v1, v2, v3])/var
    s_second = np.array([
        [0., 0., v13],
        [v12, 0., v23],
        [v13, v23, 0.]
    ])/var
    s_total = s_first + s_second.sum(axis=1)

    return s_first, s_total


def f_ishigami_vec(x):
    """Output of shape (2, n)."""
    res = f_ishigami(x)
    return res, res


class TestSobolIndices:

    dists = [
        stats.uniform(loc=-np.pi, scale=2*np.pi)  # type: ignore[attr-defined]
    ] * 3

    def test_sample_AB(self):
        # (d, n)
        A = np.array(
            [[1, 4, 7, 10],
             [2, 5, 8, 11],
             [3, 6, 9, 12]]
        )
        B = A + 100
        # (d, d, n)
        ref = np.array(
            [[[101, 104, 107, 110],
              [2, 5, 8, 11],
              [3, 6, 9, 12]],
             [[1, 4, 7, 10],
              [102, 105, 108, 111],
              [3, 6, 9, 12]],
             [[1, 4, 7, 10],
              [2, 5, 8, 11],
              [103, 106, 109, 112]]]
        )
        AB = sample_AB(A=A, B=B)
        assert_allclose(AB, ref)

    @pytest.mark.xfail_on_32bit("Can't create large array for test")
    @pytest.mark.parametrize(
        'func',
        [f_ishigami, pytest.param(f_ishigami_vec, marks=pytest.mark.slow)],
        ids=['scalar', 'vector']
    )
    def test_ishigami(self, ishigami_ref_indices, func):
        rng = np.random.default_rng(28631265345463262246170309650372465332)
        res = sobol_indices(
            func=func, n=4096,
            dists=self.dists,
            random_state=rng
        )

        if func.__name__ == 'f_ishigami_vec':
            ishigami_ref_indices = [
                    [ishigami_ref_indices[0], ishigami_ref_indices[0]],
                    [ishigami_ref_indices[1], ishigami_ref_indices[1]]
            ]

        assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
        assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-2)

        assert res._bootstrap_result is None
        bootstrap_res = res.bootstrap(n_resamples=99)
        assert isinstance(bootstrap_res, BootstrapSobolResult)
        assert isinstance(res._bootstrap_result, BootstrapResult)

        assert res._bootstrap_result.confidence_interval.low.shape[0] == 2
        assert res._bootstrap_result.confidence_interval.low[1].shape \
               == res.first_order.shape

        assert bootstrap_res.first_order.confidence_interval.low.shape \
               == res.first_order.shape
        assert bootstrap_res.total_order.confidence_interval.low.shape \
               == res.total_order.shape

        assert_array_less(
            bootstrap_res.first_order.confidence_interval.low, res.first_order
        )
        assert_array_less(
            res.first_order, bootstrap_res.first_order.confidence_interval.high
        )
        assert_array_less(
            bootstrap_res.total_order.confidence_interval.low, res.total_order
        )
        assert_array_less(
            res.total_order, bootstrap_res.total_order.confidence_interval.high
        )

        # call again to use previous results and change a param
        assert isinstance(
            res.bootstrap(confidence_level=0.9, n_resamples=99),
            BootstrapSobolResult
        )
        assert isinstance(res._bootstrap_result, BootstrapResult)

    def test_func_dict(self, ishigami_ref_indices):
        rng = np.random.default_rng(28631265345463262246170309650372465332)
        n = 4096
        dists = [
            stats.uniform(loc=-np.pi, scale=2*np.pi),
            stats.uniform(loc=-np.pi, scale=2*np.pi),
            stats.uniform(loc=-np.pi, scale=2*np.pi)
        ]

        A, B = sample_A_B(n=n, dists=dists, random_state=rng)
        AB = sample_AB(A=A, B=B)

        func = {
            'f_A': f_ishigami(A).reshape(1, -1),
            'f_B': f_ishigami(B).reshape(1, -1),
            'f_AB': f_ishigami(AB).reshape((3, 1, -1))
        }

        res = sobol_indices(
            func=func, n=n,
            dists=dists,
            random_state=rng
        )
        assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)

        res = sobol_indices(
            func=func, n=n,
            random_state=rng
        )
        assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)

    def test_method(self, ishigami_ref_indices):
        def jansen_sobol(f_A, f_B, f_AB):
            """Jansen for S and Sobol' for St.

            From Saltelli2010, table 2 formulations (c) and (e)."""
            var = np.var([f_A, f_B], axis=(0, -1))

            s = (var - 0.5*np.mean((f_B - f_AB)**2, axis=-1)) / var
            st = np.mean(f_A*(f_A - f_AB), axis=-1) / var

            return s.T, st.T

        rng = np.random.default_rng(28631265345463262246170309650372465332)
        res = sobol_indices(
            func=f_ishigami, n=4096,
            dists=self.dists,
            method=jansen_sobol,
            random_state=rng
        )

        assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
        assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-2)

        def jansen_sobol_typed(
            f_A: np.ndarray, f_B: np.ndarray, f_AB: np.ndarray
        ) -> tuple[np.ndarray, np.ndarray]:
            return jansen_sobol(f_A, f_B, f_AB)

        _ = sobol_indices(
            func=f_ishigami, n=8,
            dists=self.dists,
            method=jansen_sobol_typed,
            random_state=rng
        )

    def test_normalization(self, ishigami_ref_indices):
        rng = np.random.default_rng(28631265345463262246170309650372465332)
        res = sobol_indices(
            func=lambda x: f_ishigami(x) + 1000, n=4096,
            dists=self.dists,
            random_state=rng
        )

        assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-2)
        assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-2)

    def test_constant_function(self, ishigami_ref_indices):

        def f_ishigami_vec_const(x):
            """Output of shape (3, n)."""
            res = f_ishigami(x)
            return res, res * 0 + 10, res

        rng = np.random.default_rng(28631265345463262246170309650372465332)
        res = sobol_indices(
            func=f_ishigami_vec_const, n=4096,
            dists=self.dists,
            random_state=rng
        )

        ishigami_vec_indices = [
                [ishigami_ref_indices[0], [0, 0, 0], ishigami_ref_indices[0]],
                [ishigami_ref_indices[1], [0, 0, 0], ishigami_ref_indices[1]]
        ]

        assert_allclose(res.first_order, ishigami_vec_indices[0], atol=1e-2)
        assert_allclose(res.total_order, ishigami_vec_indices[1], atol=1e-2)

    @pytest.mark.xfail_on_32bit("Can't create large array for test")
    def test_more_converged(self, ishigami_ref_indices):
        rng = np.random.default_rng(28631265345463262246170309650372465332)
        res = sobol_indices(
            func=f_ishigami, n=2**19,  # 524288
            dists=self.dists,
            random_state=rng
        )

        assert_allclose(res.first_order, ishigami_ref_indices[0], atol=1e-4)
        assert_allclose(res.total_order, ishigami_ref_indices[1], atol=1e-4)

    def test_raises(self):

        message = r"Each distribution in `dists` must have method `ppf`"
        with pytest.raises(ValueError, match=message):
            sobol_indices(n=0, func=f_ishigami, dists="uniform")

        with pytest.raises(ValueError, match=message):
            sobol_indices(n=0, func=f_ishigami, dists=[lambda x: x])

        message = r"The balance properties of Sobol'"
        with pytest.raises(ValueError, match=message):
            sobol_indices(n=7, func=f_ishigami, dists=[stats.uniform()])

        with pytest.raises(ValueError, match=message):
            sobol_indices(n=4.1, func=f_ishigami, dists=[stats.uniform()])

        message = r"'toto' is not a valid 'method'"
        with pytest.raises(ValueError, match=message):
            sobol_indices(n=0, func=f_ishigami, method='toto')

        message = r"must have the following signature"
        with pytest.raises(ValueError, match=message):
            sobol_indices(n=0, func=f_ishigami, method=lambda x: x)

        message = r"'dists' must be defined when 'func' is a callable"
        with pytest.raises(ValueError, match=message):
            sobol_indices(n=0, func=f_ishigami)

        def func_wrong_shape_output(x):
            return x.reshape(-1, 1)

        message = r"'func' output should have a shape"
        with pytest.raises(ValueError, match=message):
            sobol_indices(
                n=2, func=func_wrong_shape_output, dists=[stats.uniform()]
            )

        message = r"When 'func' is a dictionary"
        with pytest.raises(ValueError, match=message):
            sobol_indices(
                n=2, func={'f_A': [], 'f_AB': []}, dists=[stats.uniform()]
            )

        with pytest.raises(ValueError, match=message):
            # f_B malformed
            sobol_indices(
                n=2,
                func={'f_A': [1, 2], 'f_B': [3], 'f_AB': [5, 6, 7, 8]},
            )

        with pytest.raises(ValueError, match=message):
            # f_AB malformed
            sobol_indices(
                n=2,
                func={'f_A': [1, 2], 'f_B': [3, 4], 'f_AB': [5, 6, 7]},
            )