"""Created on Jul 18 00:35:26 2024"""

from inspect import isfunction
from typing import Callable, Optional, Dict

import numpy as np
from custom_inherit import doc_inherit  # type: ignore

from .. import (
    ARC_SINE,
    BETA,
    CHI_SQUARE,
    distributions as dist,
    doc_style,
    EXPONENTIAL,
    FOLDED_NORMAL,
    GAMMA,
    GAUSSIAN,
    HALF_NORMAL,
    LAPLACE,
    LINE,
    LOG_NORMAL,
    SKEW_NORMAL,
    Params_,
    ListOrNdArray,
)

model_map = {
    ARC_SINE: dist.ArcSineDistribution,
    BETA: dist.BetaDistribution,
    CHI_SQUARE: dist.ChiSquareDistribution,
    EXPONENTIAL: dist.ExponentialDistribution,
    FOLDED_NORMAL: dist.FoldedNormalDistribution,
    GAMMA: dist.GammaDistribution,
    GAUSSIAN: dist.GaussianDistribution,
    HALF_NORMAL: dist.HalfNormalDistribution,
    LAPLACE: dist.LaplaceDistribution,
    LOG_NORMAL: dist.LogNormalDistribution,
    SKEW_NORMAL: dist.SkewNormalDistribution,
    LINE: dist.LineFunction,
}


def multi_base(
    x: ListOrNdArray,
    distribution_func: Callable,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    """Generate data based on a combination of distributions with optional noise.

    Parameters
    ----------
    x : Union[List[int | float], np.ndarray]
        Input array of values.
    distribution_func : Callable
        The distribution function to be used to generate data.
    params : Union[List[Tuple[int | float, ...]], np.ndarray]
        List of tuples containing the parameters for the required distribution.
    noise_level : float, optional
        Standard deviation of the noise to be added to the data.
        Defaults to 0.0.
    normalize : bool, optional
        If ``True``, the distribution is normalized so that the total area under the PDF equals 1.
        Defaults to ``False``.

    Returns
    -------
    np.array
        Array of the same shape as :math:`x`, containing the evaluated values.
    """
    y = np.zeros_like(a=x, dtype=float)

    for param_set in params:
        if isinstance(param_set, float):
            param_set = [param_set]
        y += distribution_func(*param_set, normalize=normalize).pdf(x)

    noise = noise_level * np.random.normal(size=y.size)
    y += noise.astype(float)

    return y


def multi_chi_squared(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""
    Generate multi-:class:`~pymultifit.distributions.chiSquare_d.ChiSquareDistribution` data with optional noise.

    Parameters
    ----------
    x : Union[List[int | float], np.ndarray]
        Input array of values.
    params : Union[List[Tuple[int | float, ...]], np.ndarray]
        List of tuples or numpy array containing the parameters for the required distribution.
    noise_level : float, optional
        Standard deviation of the noise to be added to the data.
        Defaults to 0.0.
    normalize : bool, optional
        If ``True``, the distribution is normalized so that the total area under the PDF equals 1.
        Defaults to ``False``.

    Returns
    -------
    np.array
        Array of the same shape as :math:`x`, containing the evaluated values.
    """
    return multi_base(
        x,
        distribution_func=dist.ChiSquareDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_gamma(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.gamma_d.GammaDistribution` data with optional noise."""
    return multi_base(
        x,
        distribution_func=dist.GammaDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_exponential(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.exponential_d.ExponentialDistribution` data with optional
    noise."""
    return multi_base(
        x,
        distribution_func=dist.ExponentialDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_folded_normal(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.foldedNormal_d.FoldedNormalDistribution` data with optional
    noise."""
    return multi_base(
        x,
        distribution_func=dist.FoldedNormalDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_gaussian(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.gaussian_d.GaussianDistribution` data with optional noise."""
    return multi_base(
        x,
        distribution_func=dist.GaussianDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_half_normal(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.halfNormal_d.HalfNormalDistribution` data with optional
    noise."""
    return multi_base(
        x,
        distribution_func=dist.HalfNormalDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_laplace(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.laplace_d.LaplaceDistribution` data with optional noise."""
    return multi_base(
        x,
        distribution_func=dist.LaplaceDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_log_normal(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.logNormal_d.LogNormalDistribution` data with optional noise."""
    return multi_base(
        x,
        distribution_func=dist.LogNormalDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


@doc_inherit(parent=multi_chi_squared, style=doc_style)
def multi_skew_normal(
    x: ListOrNdArray,
    params: Params_,
    noise_level: float = 0.0,
    normalize: bool = False,
) -> np.ndarray:
    r"""Generate multi-:class:`~pymultifit.distributions.skewNormal_d.SkewNormalDistribution` data with optional
    noise."""
    return multi_base(
        x,
        distribution_func=dist.SkewNormalDistribution,
        params=params,
        noise_level=noise_level,
        normalize=normalize,
    )


def multiple_models(
    x: ListOrNdArray,
    params: Params_,
    model_list: list[str],
    noise_level: float = 0.0,
    normalize: bool = False,
    mapping_dict: Optional[Dict[str, Callable]] = None,
) -> np.ndarray:
    """
    Generate data based on a combination of different models with optional noise.

    Parameters
    ----------
    x : Union[List[int | float], np.ndarray]
        Input array of values.
    params : Union[List[Tuple[int | float, ...]], np.ndarray]
        List of tuples containing the parameters for each model.
    model_list : list
        A list of model names corresponding to the models to be used.
    noise_level : float, optional
        Standard deviation of the noise to be added to the data, by default 0.0.
    normalize : bool, optional
        If ``True``, the distribution is normalized so that the total area under the PDF equals 1.
        Defaults to ``False``.
    mapping_dict: dict, optional
        A dictionary mapping between distribution names and their corresponding classes.

    Returns
    -------
    np.array
        Array of the same shape as :math:`x`, containing the evaluated values.
    """
    y = np.zeros_like(a=x, dtype=float)

    model_mapping = model_map if mapping_dict is None else mapping_dict

    for par_index, model in enumerate(model_list):
        if model in model_mapping:
            _instance = model_mapping[model]
            if isfunction(_instance):
                y += _instance(x, *params[par_index])
            else:
                y += _instance(*params[par_index], normalize=normalize).pdf(np.asarray(x))

    noise = noise_level * np.random.normal(size=y.size)
    y += noise.astype(float)

    return y
