"""
This module contains the implementation of the meta in UQTestFuns.

The implementation is based on the meta proposed in [1].
To further specify the meta, a set of default parameter values
according to [1] is also implemented.

References
----------

1. William Becker, “Metafunctions for benchmarking in sensitivity analysis,”
   Reliability Engineering & System Safety, vol. 204, p. 107189, 2020,
   doi:
   `10.1016/j.ress.2020.107189 <https://doi.org/10.1016/j.ress.2020.107189>`_.
"""
import math
import numpy as np

from dataclasses import dataclass
from numpy.typing import ArrayLike
from typing import Optional, Union, List

from .metaspec import UQMetaFunSpec, UQTestFunSpec
from .basis_functions import BASIS_BY_ID
from ..core import UQTestFun, ProbInput, UnivDist


__all__ = ["UQMetaTestFun", "default_coeffs_gen"]


def default_coeffs_gen(sample_size: int) -> np.ndarray:
    r"""Generate coefficient values for each effect term using the default.

    The default coefficient values generator is a Gaussian mixture distribution
    as proposed in [1] with the following form:

    .. math::
       \mathcal{GM}(\phi, \mu_1, \sigma_1, \mu_2, \sigma_2) =
       \phi \mathcal{N}(\mu_1, \sigma_1)
       + (1 - \phi) \mathcal{N}(\mu_2, \sigma_2)

    where :math:`\phi = 0.7`, :math:`\mu_1 = 0`, :math:`\mu_2 = 0`,
    :math:`\sigma_1 = \sqrt{0.5}`, and :math:`\sigma_2 = \sqrt{5}`.

    Parameters
    ----------
    sample_size : int
        Number of sample points to generate.

    Returns
    -------
    np.ndarray
        The coefficient values generated by the generator.
    """
    yy = np.zeros(sample_size)

    phi = 0.7
    sigma_1 = np.sqrt(0.5)
    sigma_2 = np.sqrt(5)

    choices = np.random.choice([0, 1], p=[phi, 1 - phi], size=sample_size)
    yy[choices == 0] = np.random.normal(0, sigma_1, size=np.sum(choices == 0))
    yy[choices == 1] = np.random.normal(0, sigma_2, size=np.sum(choices == 1))

    return yy


def _evaluate_test_function(
    xx: np.ndarray, uqtestfun_spec: UQTestFunSpec
) -> np.ndarray:
    """Alternative way to evaluate metafunction realizations."""

    basis_functions = uqtestfun_spec.basis_functions
    selected_basis = uqtestfun_spec.selected_basis

    # Evaluate the selected basis functions on xx
    basis_vals = np.zeros((xx.shape[0], uqtestfun_spec.spatial_dimension))
    for idx, basis_id in enumerate(selected_basis):
        basis_vals[:, idx] = basis_functions[basis_id](xx[:, idx])

    effects_tuples = uqtestfun_spec.effects_tuples
    effects_coeffs = uqtestfun_spec.effects_coeffs

    # Evaluate each effect term of the test function
    yy = np.zeros(xx.shape[0])
    for key in effects_tuples:
        n_way = effects_tuples[key]
        coeffs = effects_coeffs[key]

        for dim_indices, coeff in zip(n_way, coeffs):
            yy[:] += coeff * np.prod(basis_vals[:, dim_indices], axis=1)

    return yy


@dataclass
class UQMetaTestFun:
    """Implementation of the metafunction.

    A metafunction is a test-function-generating function.
    A random realization of a test function can be generated
    from a metafunction.

    Parameters
    ----------
    metafun_spec : UQMetaFunSpec
        An instance of meta specification that fully defines a
        meta.
    """

    metafun_spec: UQMetaFunSpec

    def get_sample(
        self, sample_size=1
    ) -> Optional[Union[UQTestFun, List[UQTestFun]]]:
        """Generate realizations of test function.

        Parameters
        ----------
        sample_size : int
            Number of test function realizations.

        Returns
        -------
        Optional[Union[UQTestFun, List[UQTestFun]]]
            Realization(s) of UQTestFunSpec.

        Notes
        -----
        - With sample_size larger than one, a list of UQTestFun instances
          is returned.
        """
        if sample_size < 1:
            return None

        # Generate realizations of test function specification
        testfun_specs = self.metafun_spec.get_sample(sample_size)

        # Create the name
        name = "metafunction_realization"
        # Create the evaluation callable
        evaluate = _evaluate_test_function

        if sample_size == 1:
            assert testfun_specs is not None
            assert not isinstance(testfun_specs, list)
            # Create an instance of inputs
            prob_input = ProbInput(testfun_specs.inputs)
            # Assign the realized spec as a parameter
            parameters = testfun_specs

            return UQTestFun(
                evaluate=evaluate,
                prob_input=prob_input,
                parameters=parameters,
                name=name,
            )

        sample = []
        for i in range(sample_size):
            # Create an instance of inputs
            assert isinstance(testfun_specs, list)
            prob_input = ProbInput(testfun_specs[i].inputs)
            # Assign the realized spec as a parameter
            parameters = testfun_specs[i]

            sample.append(
                UQTestFun(
                    evaluate=evaluate,
                    prob_input=prob_input,
                    parameters=parameters,
                    name=name,
                )
            )

        return sample

    @classmethod
    def from_default(
        cls,
        spatial_dimension: ArrayLike,
        input_id: Optional[int] = None,
    ):
        """Create a metafunction with parameters according to Becker (2019).

        Parameters
        ----------
        spatial_dimension : Union[int, Sized]
            Number of dimensions of the test functions generated
            by the meta. If a set of values are given, a single value
            will be selected at random.

        Returns
        -------
        UQMetaTestFun
            An instance of metafunction with the default parameters.
        """
        # Select a single spatial dimension randomly
        if not isinstance(spatial_dimension, int):
            spatial_dimension = np.asarray(spatial_dimension)
            spatial_dimension = int(np.random.choice(spatial_dimension))

        effects_dict = {
            1: None,  # Take all
            2: math.floor(0.5 * spatial_dimension),
            3: math.floor(0.2 * spatial_dimension),
        }

        if input_id is None:
            input_id = np.random.randint(0, 8)

        input_marginals = [
            UnivDist(distribution="uniform", parameters=[0, 1]),
            UnivDist(
                distribution="trunc-normal",
                parameters=[0.5, 0.15, 0.0, 1.0],
            ),
            UnivDist(distribution="beta", parameters=[8.0, 2.0, 0.0, 1.0]),
            UnivDist(distribution="beta", parameters=[2.0, 8.0, 0.0, 1.0]),
            UnivDist(distribution="beta", parameters=[2.0, 0.8, 0.0, 1.0]),
            UnivDist(distribution="beta", parameters=[0.8, 2.0, 0.0, 1.0]),
            UnivDist(distribution="logitnormal", parameters=[0.0, 3.16]),
        ]

        if input_id < 7:
            input_marginals = [input_marginals[input_id]]

        metafun_spec = UQMetaFunSpec(
            spatial_dimension=spatial_dimension,
            basis_functions=BASIS_BY_ID,
            effects_dict=effects_dict,
            input_marginals=input_marginals,
            coeffs_generator=default_coeffs_gen,
        )

        return cls(metafun_spec)
