import typing as t

from sarus_statistics.ops.histograms.local import private_histogram
from scipy import stats
import numpy as np
import pandas as pd


def update_quantiles(
    probabilities: t.List[float],
    quantiles: t.List[float],
    is_int: bool,
) -> t.Tuple[t.List[float], t.List[float]]:
    probabilities_array = np.array(sorted(probabilities))
    quantiles_array = np.array(sorted(quantiles))
    if is_int:
        quantiles_array = np.round(quantiles_array).astype(int)
    _, first_index = np.unique(quantiles_array, return_index=True)
    _, last_index = np.unique(np.flip(quantiles_array), return_index=True)
    last_index = len(quantiles_array) - last_index - 1
    total_indices = np.unique(np.concatenate([first_index, last_index]))
    final_quantiles = quantiles_array[total_indices]
    probabilities_array = probabilities_array[total_indices]
    return probabilities_array.tolist(), final_quantiles.tolist()


def quantiles_values_from_distribution(
    distribution_name: str,
    params: t.Dict[str, float],
    nb_quantiles: int,
    bounds: t.Tuple[float, float],
) -> t.List[float]:
    """It returns the quatile values for specific distribution for then
    computing private_histograms. Endpoints are not included.

    Args:
        distribution_name (str): name of the distribution
        params (Dict[str, float]): parameters
        nb_quantiles (int): number of quantiles to estimate
        bounds (Tuple[float, float]): lower and upper bound of data

    Raises:
        NotImplementedError: if the given distribution is not implemented

    Returns:
        List[float]: quantiles values.
    """
    proba = np.linspace(0, 1, nb_quantiles + 1, endpoint=False)[1:]

    if distribution_name == "Uniform":
        quantiles = stats.uniform.ppf(
            proba, loc=bounds[0], scale=bounds[1] - bounds[0]
        )

    elif distribution_name == "Normal":
        mu = params["mu"]
        sigma = params["sigma"]
        quantiles = stats.norm.ppf(proba, mu, sigma)

    elif distribution_name == "Exponential":
        lambd = params["lambda"]
        quantiles = stats.expon.ppf(proba, loc=bounds[0], scale=1 / lambd)

    elif distribution_name == "Beta":
        alpha = params["alpha"]
        beta = params["beta"]
        quantiles = stats.beta.ppf(
            proba, alpha, beta, loc=bounds[0], scale=bounds[1]
        )

    elif distribution_name == "Gamma":
        k = params["k"]
        theta = params["theta"]
        quantiles = stats.gamma.ppf(proba, k, loc=bounds[0], scale=theta)

    elif distribution_name == "Pareto":
        alpha = params["alpha"]
        scale = params["scale"]
        quantiles = stats.pareto.ppf(
            proba, alpha, loc=bounds[0] - 1.0, scale=scale
        )
    else:
        raise NotImplementedError(
            f"{distribution_name} distribution model not yet implemented"
        )

    return t.cast(t.List[float], quantiles.tolist())


def private_quantiles_from_distribution(
    data: pd.DataFrame,
    data_col: str,
    user_col: str,
    private_col: str,
    weight_col: str,
    noise: float,
    nb_quantiles: int,
    bounds: t.Tuple[float, float],
    max_multiplicity: float,
    distribution_name: str,
    params: t.Dict[str, float],
    random_generator: t.Optional[np.random.Generator] = None,
) -> t.Dict[float, float]:
    """Compute DP quantiles using private histogram query.
    It computes bins (quantile values) from the inferred distribution,
    it preprocess data[user_col] to be categorical using the bins,
    It computes private histograms and then the probabilities


    Args:
        data (pd.DataFrame): dataset
        data_col (str): colum to be evaluated
        user_col (str):
        private_col (str):
        weight_col (str):
        noise (float):
        nb_quantiles (int): number of quantiles
        bounds (t.Tuple[float, float]):
        max_multiplicity (float):
        distribution_name (str): e.g. "Uniform", "Normal" etc
        params (t.Dict[str, float]): specific to the distribution
        random_generator (t.Optional[np.random.Generator], optional):
            Defaults to None.

    Returns:
        t.Dict[float, float]: private quantiles
    """

    quantiles_values = quantiles_values_from_distribution(
        nb_quantiles=nb_quantiles,
        bounds=bounds,
        distribution_name=distribution_name,
        params=params,
    )

    # sorted bins
    bins = [bounds[0]] + list(quantiles_values) + [bounds[1]]

    # if equal values in bins, add small noise, add noise and resize to bounds
    if len(bins) != len(set(bins)):
        rho = 1e-8  # noise
        bins = [b + np.random.uniform(-rho, rho) for b in bins]
        bins = sorted(bins)
        bins = [
            (b - min(bins)) / max(bins) * (bounds[1] - bounds[0]) + bounds[0]
            for b in bins
        ]

    # replacing data[data_col] with the associated categorical bin
    data[data_col] = pd.cut(
        data[data_col], bins=np.asarray(bins), include_lowest=True
    )

    hist_dict = private_histogram(
        data=data,
        data_col=data_col,
        user_col=user_col,
        private_col=private_col,
        weight_col=weight_col,
        noise=noise,
        max_multiplicity=max_multiplicity,
        random_generator=random_generator,
    )

    # it needs to be sorted by keys (bins)
    histogram_values: np.ndarray = np.asarray(
        [0.0] + [x[1] for x in sorted(hist_dict.items())]
    )

    proba = np.cumsum(histogram_values / np.sum(histogram_values)).tolist()

    # numerical errors make the last value deviate lightly from 1.0,
    # here forcing it to be exactly 1.0
    proba[-1] = 1.0
    return dict(zip(proba, bins))