from typing import Optional, cast

import numpy as np

from sarus_statistics.ops.utils import generator_from_seed

# def smooth_utility(data, probs, rho):
#    """add smoothing to probabilities for the exponential mechanism"""
#    split = len(data) // 2
#    if len(data) % 2 == 0:
#        w += (-rho/2) * tf.concat([
#            tf.ones(
#                (*w.shape[:-1], split)
#            ),
#            tf.zeros(
#                (*w.shape[:-1], split)
#            )
#    ], axis=-1) + (rho/2) * tf.concat([
#        tf.zeros(
#            (*w.shape[:-1], split)
#        ),
#        tf.ones(
#            (*w.shape[:-1], split)
#        )
#    ], axis=-1)
#    else:
#        w = tf.concat(
#            [w[..., :split + 1] - (rho/2), w[..., split:] + (rho/2)],
#            -1
#        )
#        probs.insert(split, probs[split])
#    return w, probs


def exp_quantile(
    sorted_data: np.ndarray,
    data_low: float,
    data_high: float,
    quantile: float,
    eps: float,
    rho: float = 1e-3,
    random_generator: Optional[np.random.Generator] = None,
) -> float:
    random_generator = (
        random_generator
        if random_generator is not None
        else generator_from_seed(random_generator)
    )
    n = len(sorted_data)
    sorted_data = np.sort(
        np.clip(
            sorted_data
            + random_generator.uniform(
                -rho * (data_high - data_low),
                rho * (data_high - data_low),
                sorted_data.shape,
            ),
            data_low,
            data_high,
        )
    )  # replace by true smoothing
    data = np.block([data_low, sorted_data, data_high])
    exp_probs = np.squeeze(
        np.array(
            [
                np.exp(-eps / 4.0 * np.abs(i - quantile * n))
                for i in range(n + 1)
            ]
        )
    )

    # compute pdf integrated over each interval
    # TODO: use logs to make it more computationnally stable
    intervals = data[1:] - data[:-1]
    pdf = np.multiply(exp_probs, intervals)
    pdf /= np.sum(pdf)
    cdf = np.cumsum(pdf)

    index = np.searchsorted(cdf, random_generator.uniform(), side='right')
    return cast(float, random_generator.uniform(data[index], data[index + 1]))