"""Functions related to parameter sampling"""

from typing import Sequence, Tuple

import numpy as np
import pandas as pd

from . import parameters
from .C import *  # noqa: F403

__all__ = ["sample_from_prior", "sample_parameter_startpoints"]


def sample_from_prior(
    prior: Tuple[str, list, str, list], n_starts: int
) -> np.array:
    """Creates samples for one parameter based on prior

    Arguments:
        prior: A tuple as obtained from
            :func:`petab.parameter.get_priors_from_df`
        n_starts: Number of samples

    Returns:
        Array with sampled values
    """
    # unpack info
    p_type, p_params, scaling, bounds = prior

    # define a function to rescale the sampled points to parameter scale
    def scale(x):
        if scaling == LIN:
            return x
        if scaling == LOG:
            return np.log(x)
        if scaling == LOG10:
            return np.log10(x)
        raise NotImplementedError(
            f"Parameter priors on the parameter scale {scaling} are "
            "currently not implemented."
        )

    def clip_to_bounds(x: np.array):
        """Clip values in array x to bounds"""
        return np.maximum(np.minimum(scale(bounds[1]), x), scale(bounds[0]))

    # define lambda functions for each parameter
    if p_type == UNIFORM:
        sp = scale(
            (p_params[1] - p_params[0]) * np.random.random((n_starts,))
            + p_params[0]
        )

    elif p_type == PARAMETER_SCALE_UNIFORM:
        sp = (p_params[1] - p_params[0]) * np.random.random(
            (n_starts,)
        ) + p_params[0]

    elif p_type == NORMAL:
        sp = scale(
            np.random.normal(
                loc=p_params[0], scale=p_params[1], size=(n_starts,)
            )
        )

    elif p_type == LOG_NORMAL:
        sp = scale(
            np.exp(
                np.random.normal(
                    loc=p_params[0], scale=p_params[1], size=(n_starts,)
                )
            )
        )

    elif p_type == PARAMETER_SCALE_NORMAL:
        sp = np.random.normal(
            loc=p_params[0], scale=p_params[1], size=(n_starts,)
        )

    elif p_type == LAPLACE:
        sp = scale(
            np.random.laplace(
                loc=p_params[0], scale=p_params[1], size=(n_starts,)
            )
        )

    elif p_type == LOG_LAPLACE:
        sp = scale(
            np.exp(
                np.random.laplace(
                    loc=p_params[0], scale=p_params[1], size=(n_starts,)
                )
            )
        )

    elif p_type == PARAMETER_SCALE_LAPLACE:
        sp = np.random.laplace(
            loc=p_params[0], scale=p_params[1], size=(n_starts,)
        )

    else:
        raise NotImplementedError(
            f"Parameter priors of type {prior[0]} are not implemented."
        )

    return clip_to_bounds(sp)


def sample_parameter_startpoints(
    parameter_df: pd.DataFrame,
    n_starts: int = 100,
    seed: int = None,
    parameter_ids: Sequence[str] = None,
) -> np.array:
    """Create :class:`numpy.array` with starting points for an optimization

    Arguments:
        parameter_df: PEtab parameter DataFrame
        n_starts: Number of points to be sampled
        seed: Random number generator seed (see :func:`numpy.random.seed`)
        parameter_ids: A sequence of parameter IDs for which to sample starting points.
            For subsetting or reordering the parameters.
            Defaults to all estimated parameters.

    Returns:
        Array of sampled starting points with dimensions
        `n_startpoints` x `n_optimization_parameters`
    """
    if seed is not None:
        np.random.seed(seed)

    # get types and parameters of priors from dataframe
    prior_list = parameters.get_priors_from_df(
        parameter_df, mode=INITIALIZATION, parameter_ids=parameter_ids
    )

    startpoints = [sample_from_prior(prior, n_starts) for prior in prior_list]

    return np.array(startpoints).T
