import numpy as np
import pandas as pd
from pymasq.errors import DataTypeError, InputError
from SALib.sample import saltelli
from SALib.analyze import sobol
from typing import Dict, Optional, Tuple, Final

__all__ = ["gen_params", "gen_samples", "analyze", "UNIFORM", "SAMPLE"]


UNIFORM: Final = "uniform"
SAMPLE: Final = "sample"


def gen_params(data: pd.DataFrame) -> Dict:
    """
    Returns a dictionary containing the metadata required for SA analysis

    Parameters
    ----------
    data : pd.DataFrame
        The input data frame. All columns must be numeric in order to perform sensitivity analysis

    Returns
    -------
    Dict
        The dictionary containing the number of columns (`num_cols`), column names (`names`),
        and min and max bounds for each column (`bounds`)
    """
    if not len(data.describe().columns) == len(
        data.columns
    ):  # sa requires a numeric dataframe
        raise DataTypeError(
            "all column dtypes must be numeric, received: \n{}".format(data.dtypes)
        )
    cols = data.columns
    min_max_vals = [[np.min(data[col]), np.max(data[col])] for col in cols]
    params = {
        "num_vars": len(cols),
        "names": cols.to_list(),
        "bounds": min_max_vals,
    }
    return params


def gen_samples(
    data: pd.DataFrame,
    method: Optional[str] = None,
    num_samples: int = 10000,
    calc_second_order: bool = True,
    params: Optional[Dict] = None,
) -> np.ndarray:
    """
    A function for generating input samples to the analyze method.

    With the `uniform` method, inputs will be generated based on a uniform distribution over the
    range of values that the input data can have. The number of samples returned by the function
    will be 𝑁*(2𝐷+2) for `calc_second_order=True` and 𝑁*(𝐷+2) for `calc_second_order=False` where
    𝑁 is the value supplied by `num_samples` and 𝐷 is the number of columns in the data.

    TODO: sample method description once implemented

    Parameters
    ----------
    data : pd.DataFrame
        A dataframe of input values to the model that sensitivity analysis should be performed on
    method : {"uniform", "sample"}, optional
        A string that dictates which form of sample generation to use. (Default: "uniform")
    num_samples : int, optional
        An integer used in computing the number of samples to return. It is NOT the returned
        number of samples but a value used to generate an appropriate number of samples
        (described  above). (Default: 10000)
    calc_second_order : bool, optional
        A boolean that determines whether the number of generated samples should conform to the
        number required for first AND second order sensitivities, or just first order
        sensitivities (as described above). (Default: True)
    params : dict or None
        The params dictionary generated by `gen_params`. If the method is "uniform" and this value
        is None, `gen_params` will be called. (Default: None)

    Returns
    -------
    np.ndarray
        An n-dimensional numpy array containing a number of samples required for first or
        second order sensitivity analysis.

    Raises
    ------
    InputError
        Throws an error if the user tries to enter a method that is not "uniform" or "sample"
    """
    if method is None or method == UNIFORM:
        if params is None:
            params = gen_params(data)
        samples = saltelli.sample(
            params, num_samples, calc_second_order=calc_second_order
        )
    elif method == SAMPLE:
        # TODO: figure out sampling method
        samples = None
    else:
        raise InputError(
            f"'method' parameter must be in ['{UNIFORM}', '{SAMPLE}'], entered '{method}'"
        )
    return samples


def analyze(
    model,
    params: Dict,
    samples: np.ndarray,
    confidence: float = 0.95,
    calc_second_order: bool = True,
) -> Tuple[np.array, np.array, Optional[np.ndarray]]:
    """
    Analyze takes in a trained model and a params dictionary describing the model inputs. The
    function calculates the first order sensitivities and total sensitivities for the inputs
    described in the params file and will optionally do the same if `calc_second_order` is `True` (default behavior).

    Parameters
    ----------
    model : child class of sklearn.base.BaseEstimator
        The trained model that will be the target for sensitivity analysis. Must be able to run
        `.predict()` on the provided set of samples
    params : Dict
        A params dict generated by the `gen_params` function
    samples : np.ndarray
        A samples array generated by the `gen_samples` function
    confidence : float [0.0, 1.0), optional
        The confidence interval provided as a value between 0 and 1. (Default: 0.95)
    calc_second_order : bool, optional
        A boolean that determines whether the number of generated samples should conform to the
        number required for first AND second order sensitivities, or just first order
        sensitivities. This boolean value should match that of the flag provided to the
        `gen_samples` function used to generate the `samples` parameter. (Default: True)

    Returns
    -------
    np.array
        An array of pd.Interval objects for the first order sensitivities and their confidence
        intervals
    np.array
        An array of pd.Interval objects for the total sensitivities and their confidence
        intervals
    np.ndarray
        A two dimensional array of pd.Interval objects for the second order sensitivities and their
        confidence intervals. Values will only exist in the top-right half of the matrix.

    See Also
    --------
    gen_params
    gen_samples
    """

    y_pred = model.predict(samples)
    Si = sobol.analyze(
        params, y_pred, calc_second_order=calc_second_order, conf_level=confidence
    )
    s1 = np.array(
        [
            pd.Interval(mean - conf / 2, mean + conf / 2, "both")
            for mean, conf in zip(Si["S1"], Si["S1_conf"])
        ]
    )
    st = np.array(
        [
            pd.Interval(mean - conf / 2, mean + conf / 2, "both")
            for mean, conf in zip(Si["ST"], Si["ST_conf"])
        ]
    )
    if calc_second_order:
        s2 = np.array(
            [
                [
                    pd.Interval(mean - conf / 2, mean + conf / 2, "both")
                    if not np.isnan(mean)
                    else np.nan
                    for mean, conf in zip(Si["S2"][row], Si["S2_conf"][row])
                ]
                for row in range(len(Si["S2"]))
            ]
        )
    else:
        s2 = None
    return s1, st, s2
