"""Measurements on Pandas Series."""

# TODO(#792): Add link to open-source paper.
# TODO(#693): Check edge cases for aggregations.
# TODO(#1023): Handle clamping bounds approximation.

# SPDX-License-Identifier: Apache-2.0
# Copyright Tumult Labs 2024

import math
from abc import abstractmethod
from itertools import accumulate
from typing import Any, List, NamedTuple, Tuple, Union, cast

import numpy as np
import pandas as pd
from pyspark.sql.types import DataType, DoubleType, LongType
from typeguard import typechecked

from tmlt.core.domains.numpy_domains import (
    NumpyDomain,
    NumpyFloatDomain,
    NumpyIntegerDomain,
)
from tmlt.core.domains.pandas_domains import PandasSeriesDomain
from tmlt.core.exceptions import UnsupportedDomainError
from tmlt.core.measurements.base import Measurement
from tmlt.core.measurements.noise_mechanisms import (
    AddDiscreteGaussianNoise,
    AddGaussianNoise,
    AddGeometricNoise,
    AddLaplaceNoise,
)
from tmlt.core.measures import Measure, PureDP, RhoZCDP
from tmlt.core.metrics import (
    AbsoluteDifference,
    HammingDistance,
    RootSumOfSquared,
    SumOf,
    SymmetricDifference,
)
from tmlt.core.random.rng import prng
from tmlt.core.random.uniform import uniform
from tmlt.core.utils.arb import (
    Arb,
    arb_abs,
    arb_add,
    arb_div,
    arb_log,
    arb_max,
    arb_mul,
    arb_sub,
)
from tmlt.core.utils.exact_number import ExactNumber, ExactNumberInput
from tmlt.core.utils.validation import validate_exact_number


class Aggregate(Measurement):
    """Aggregate a Pandas Series and produce a float or int."""

    @typechecked
    def __init__(
        self,
        input_domain: PandasSeriesDomain,
        input_metric: Union[HammingDistance, SymmetricDifference],
        output_measure: Measure,
        output_spark_type: DataType,
    ):
        """Constructor.

        Args:
            input_domain: Input domain. Must have type PandasSeriesDomain.
            input_metric: Input metric.
            output_measure: Output measure.
            output_spark_type: Spark DataType of the output. This is required to use
                this measurement within a udf.
        """
        self._output_spark_type = output_spark_type
        super().__init__(
            input_domain=input_domain,
            input_metric=input_metric,
            output_measure=output_measure,
            is_interactive=False,
        )

    @property
    def output_spark_type(self) -> DataType:
        """Return the Spark type of the aggregated value."""
        return self._output_spark_type

    @abstractmethod
    def __call__(self, data: pd.Series) -> Union[float, int]:
        """Perform measurement."""


class NoisyBounds(Aggregate):
    """Estimates the bounds of a Pandas series."""

    @typechecked
    def __init__(
        self,
        input_domain: PandasSeriesDomain,
        alpha: ExactNumberInput,
        threshold_fraction: float = 0.95,
    ):
        r"""Constructor.

        Args:
            input_domain: PandasSeriesDomain or either integers or floats.
            alpha: The noise scale parameter for Geometric noise that will be added
                to true number of values falling between the tested bounds on each
                round of the algorithm.
                Noise with scale of :math:`\alpha / 2` will be added
                to compute the threshold.
                See :class:`~.AddGeometricNoise` for more information.
            threshold_fraction: The fraction of the total count to use as the threshold.
                This value should be between (0, 1]. By default it is set to 0.95.
        """
        element_type = input_domain.element_domain
        if isinstance(element_type, NumpyFloatDomain):
            if element_type.allow_nan:
                raise UnsupportedDomainError(
                    input_domain, "Input domain must disallow NaNs."
                )
            splits = (
                [-(2**i) * 2**-100 for i in range(200, -1, -1)]
                + [0]
                + [2**i * 2**-100 for i in range(201)]
            )
        elif isinstance(element_type, NumpyIntegerDomain):
            splits = (
                [-(2 ** (element_type.size - 1)) + 1]
                + [-(2**i) for i in range(element_type.size - 2, -1, -1)]
                + [0]
                + [2**i for i in range(element_type.size)]
            )
        else:
            raise ValueError(
                "Invalid element type, expected [NumpyFloatDomain, "
                f"NumpyIntegerDomain], got {element_type}."
            )
        self._splits = splits
        self._element_type = element_type

        try:
            validate_exact_number(
                value=alpha,
                allow_nonintegral=True,
                minimum=0,
                minimum_is_inclusive=True,
            )
        except ValueError as e:
            raise ValueError(f"Invalid noise scale: {e}") from e
        if not 0 < threshold_fraction <= 1:
            raise ValueError(
                f"Invalid threshold fraction: {threshold_fraction}. Must be in (0, 1]."
            )
        self._alpha = ExactNumber(alpha)
        self._threshold_fraction = threshold_fraction
        super().__init__(
            input_domain=input_domain,
            input_metric=SymmetricDifference(),
            output_measure=PureDP(),
            output_spark_type=self.output_spark_type,
        )

    @property
    def alpha(self) -> ExactNumber:
        """Returns the alpha."""
        return self._alpha

    @property
    def threshold_fraction(self) -> float:
        """Returns the threshold."""
        return self._threshold_fraction

    @property
    def splits(self) -> Union[List[float], List[int]]:
        """Returns the splits."""
        return self._splits.copy()

    @property
    def element_type(self) -> NumpyDomain:
        """Returns the element type of the series."""
        return self._element_type

    @property
    def output_spark_type(self) -> DataType:
        """Returns the output Spark type after being used as a UDF."""
        return (
            DoubleType()
            if isinstance(self.element_type, NumpyFloatDomain)
            else LongType()
        )

    @typechecked
    def privacy_function(self, d_in: ExactNumberInput) -> ExactNumber:
        """Returns the smallest d_out satisfied by the measurement.

        See `the architecture guide
        <https://docs.tmlt.dev/core/latest/topic-guides/architecture.html>`_ for more
        information.

        Args:
            d_in: Distance between inputs under input_metric.
        """
        self.input_metric.validate(d_in)
        d_in = ExactNumber(d_in)
        if d_in == 0:
            return ExactNumber(0)
        if self.alpha == 0:
            return ExactNumber(float("inf"))
        if d_in < 1:
            raise NotImplementedError()
        return (4 / self.alpha) * d_in

    def __call__(self, data: pd.Series) -> Union[float, int]:
        """Returns the bounds for the given column."""
        # Creates a series mapping bin index to count of values in that bin. Indices are
        # such that 6 bins will result in indices -3, -2, -1, 1, 2, 3, to make matching
        # corresponding positive and negative bins easier.  This includes zeros for bins
        # with no values. If any values do not fall in bin (floats larger than 2^100 or
        # smaller than -2^100), they are dropped.
        num_bins = len(self._splits) - 1
        bin_counts = pd.cut(
            data,
            bins=self._splits,
            include_lowest=True,
            labels=list(range(-num_bins // 2, 0)) + list(range(1, num_bins // 2 + 1)),
        ).value_counts()

        cumulative_counts = list(
            accumulate(
                bin_counts[i] + bin_counts[-i] for i in range(1, num_bins // 2 + 1)
            )
        )

        add_threshold_noise = AddGeometricNoise(self._alpha / 2)
        noisy_threshold = add_threshold_noise(
            np.int64(self.threshold_fraction * cumulative_counts[-1])
        )

        # Finding the bin that contains enough counts above the threshold
        add_bin_noise = AddGeometricNoise(self.alpha)
        bin_upper_limits = self._splits[len(self._splits) // 2 + 1 :]

        for i, bin_count in enumerate(cumulative_counts):
            noisy_count = add_bin_noise(np.int64(bin_count))
            if noisy_count >= noisy_threshold:
                return bin_upper_limits[i]
        return bin_upper_limits[-1]


class NoisyQuantile(Aggregate):
    """Estimates the quantile of a Pandas Series."""

    @typechecked
    def __init__(
        self,
        input_domain: PandasSeriesDomain,
        output_measure: Union[PureDP, RhoZCDP],
        quantile: float,
        lower: Union[float, int],
        upper: Union[float, int],
        epsilon: ExactNumberInput,
    ):
        """Constructor.

        Args:
            input_domain: Input domain. Must be PandasSeriesDomain.
            output_measure: Output measure.
            quantile: The quantile to produce.
            lower: The lower clamping bound.
            upper: The upper clamping bound.
            epsilon: The pure-dp privacy parameter to use to produce the quantile.
        """
        if not 0 <= quantile <= 1:
            raise ValueError("Quantile must be between 0 and 1.")

        if math.isnan(lower) or math.isinf(lower):
            raise ValueError(
                f"Lower clamping bound must be finite and non-nan, not {lower}."
            )
        if math.isnan(upper) or math.isinf(upper):
            raise ValueError(
                f"Upper clamping bound must be finite and non-nan, not {upper}."
            )
        if lower > upper:
            raise ValueError(
                f"Lower bound ({lower}) can not be greater than "
                f"the upper bound ({upper})."
            )
        PureDP().validate(epsilon)
        if not isinstance(
            input_domain.element_domain, (NumpyIntegerDomain, NumpyFloatDomain)
        ):
            raise UnsupportedDomainError(
                input_domain,
                (
                    "input_domain.element_domain must be NumpyIntegerDomain or"
                    " NumpyFloatDomain, not"
                    f" {type(input_domain.element_domain).__name__}"
                ),
            )

        if (
            isinstance(input_domain.element_domain, NumpyFloatDomain)
            and input_domain.element_domain.allow_nan
        ):
            raise UnsupportedDomainError(
                input_domain, "Input domain must disallow NaNs."
            )

        self._quantile = quantile
        self._epsilon = ExactNumber(epsilon)
        self._lower = lower
        self._upper = upper

        super().__init__(
            input_domain=input_domain,
            input_metric=SymmetricDifference(),
            output_measure=output_measure,
            output_spark_type=DoubleType(),
        )

    @property
    def quantile(self) -> float:
        """Returns the quantile to be computed."""
        return self._quantile

    @property
    def lower(self) -> Union[float, int]:
        """Returns the lower clamping bound."""
        return self._lower

    @property
    def upper(self) -> Union[float, int]:
        """Returns the upper clamping bound."""
        return self._upper

    @property
    def epsilon(self) -> ExactNumber:
        """Returns the PureDP privacy budget to be used for producing a quantile."""
        return self._epsilon

    @typechecked
    def privacy_function(self, d_in: ExactNumberInput) -> ExactNumber:
        r"""Returns the smallest d_out satisfied by the measurement.

        This algorithm uses the exponential mechanism, so benefits from the same privacy
        analysis:

        If the output measure is :class:`~.PureDP`, returns

            :math:`\epsilon \cdot d_{in}`

        If the output measure is :class:`~.RhoZCDP`, returns

            :math:`\frac{1}{8}(\epsilon \cdot d_{in})^2`

        where:

        * :math:`d_{in}` is the input argument ``d_in``
        * :math:`\epsilon` is :attr:`~.epsilon`

        See :cite:`Cesar021` for the zCDP privacy analysis.

        Args:
            d_in: Distance between inputs under input_metric.
        """
        self.input_metric.validate(d_in)
        d_in = ExactNumber(d_in)
        if self.output_measure == PureDP():
            return self.epsilon * d_in
        assert self.output_measure == RhoZCDP()
        return (self.epsilon * d_in) ** 2 / 8

    def __call__(self, data: pd.Series) -> float:
        """Return DP answer(float) to quantile query.

        TODO(#792) Add link to open-source paper: See this document for a description
        of the algorithm.

        Args:
            data: The Series on which to compute the quantile.
        """
        if self.lower == self.upper:
            return self.lower

        float_epsilon = self.epsilon.to_float(round_up=False)

        l, u = _select_quantile_interval(
            values=data.to_list(),
            q=self.quantile,
            epsilon=float_epsilon,
            lower=self.lower,
            upper=self.upper,
        )

        # sample uniformly from bin
        return uniform(lower=l, upper=u)


class AddNoiseToSeries(Measurement):
    """A measurement that adds noise to each value in a pandas Series."""

    @typechecked
    def __init__(
        self,
        noise_measurement: Union[
            AddLaplaceNoise,
            AddGeometricNoise,
            AddDiscreteGaussianNoise,
            AddGaussianNoise,
        ],
    ):
        """Constructor.

        Args:
            noise_measurement: Noise Measurement to be applied to each element
                in input pandas Series.
        """
        if not noise_measurement.output_measure in [PureDP(), RhoZCDP()]:
            raise AssertionError(
                "This is probably a bug; please let us know so we can fix it!"
            )
        input_metric = (
            SumOf(AbsoluteDifference())
            if noise_measurement.output_measure == PureDP()
            else RootSumOfSquared(AbsoluteDifference())
        )
        self._noise_measurement = noise_measurement
        super().__init__(
            input_domain=PandasSeriesDomain(noise_measurement.input_domain),
            input_metric=input_metric,
            output_measure=noise_measurement.output_measure,
            is_interactive=False,
        )

    @property
    def noise_measurement(
        self,
    ) -> Union[
        AddLaplaceNoise, AddGeometricNoise, AddDiscreteGaussianNoise, AddGaussianNoise
    ]:
        """Returns measurement that adds noise to each number in pandas Series."""
        return self._noise_measurement

    @property
    def input_domain(self) -> PandasSeriesDomain:
        """Return input domain for the measurement."""
        return cast(PandasSeriesDomain, super().input_domain)

    @property
    def output_type(self) -> DataType:
        """Return the output data type after being used as a UDF."""
        return self.noise_measurement.output_type

    @typechecked
    def privacy_function(self, d_in: ExactNumberInput) -> ExactNumber:
        r"""Returns the smallest d_out satisfied by the measurement.

        Args:
            d_in: Distance between inputs under input_metric.
        """
        return self.noise_measurement.privacy_function(d_in)

    def __call__(self, values: pd.Series) -> pd.Series:
        """Adds noise to each number in the input Series."""
        return values.apply(
            lambda x: self.noise_measurement(x)  # pylint: disable=unnecessary-lambda
        )


class _RankedInterval(NamedTuple):
    """A interval and its rank w.r.t some list."""

    rank: Arb
    """Rank of any number in this interval."""

    lower: Arb
    """Lower endpoint of the interval."""

    upper: Arb
    """Upper endpoint of the interval."""


def _get_intervals_with_ranks(
    values: Union[List[float], "np.ndarray[Any, Any]"], lower: float, upper: float
) -> List[_RankedInterval]:
    """Returns a list of intervals constructed from ``values``.

    The list of intervals returned consist of three types of intervals:

        - One interval between ``lower`` and the smallest number in ``values``
          strictly larger than ``lower``.
        - Non-empty intervals between consecutive numbers in ``values`` (sorted in
          ascending order)
        - One interval from the largest number in ``values`` strictly smaller than
          ``upper`` and ``upper``.

    The rank associated with each interval is the rank of any number in the interval
    w.r.t all numbers in ``values``.
    """
    values = np.sort(values)
    lower_index, upper_index = np.searchsorted(values, [lower, upper], side="left")
    # if a value was inserted at lower_index, its rank would be lower_index
    # every value below lower_index is strictly below lower

    intervals = []
    left_float = lower
    left_arb = Arb.from_float(lower)
    for index in range(lower_index, upper_index):
        right_float = float(values[index])
        if left_float < right_float:
            right_arb = Arb.from_float(right_float)
            intervals.append(
                _RankedInterval(
                    rank=Arb.from_int(index), lower=left_arb, upper=right_arb
                )
            )
            left_float = right_float
            left_arb = right_arb

    intervals.append(
        _RankedInterval(
            rank=Arb.from_int(int(upper_index)),
            lower=left_arb,
            upper=Arb.from_float(upper),
        )
    )
    # Note that the `_RankedInterval`s are constructed with exact arbs (with radius=0)
    # This guarantees that the arbs always represent the floating point numbers exactly
    # regardless of what the precision is.
    return intervals


def _select_quantile_interval(
    values: List[float], q: float, epsilon: float, lower: float, upper: float
) -> Tuple[float, float]:
    r"""Returns a privately selected interval (l, u) with the max noisy score.

    In particular, this function performs the following steps:
        - Constructs a list of intervals by calling ``_get_intervals_with_ranks``.
        - Compute the target rank as: :math:`target = q * len(values)`.
        - Assigns each interval :math:`(rank, x_i, x_j)` a noisy score computed as:
            :math:`log(x_j - x_i) - |rank - target| * \frac{epsilon}{2 \cdot \Delta U} + G`
            where :math:`G` is a sampled from the standard Gumbel distribution.
        - Returns the interval with the highest noisy score.
    """  # pylint:disable=line-too-long
    arb_q = Arb.from_float(float(q))
    prec = 53
    # target_rank = arb_q * len(values)
    target_rank = arb_mul(arb_q, Arb.from_int(len(values)), prec)

    # Get bin ranks
    intervals = _get_intervals_with_ranks(values, lower, upper)

    if epsilon == float("inf"):
        intervals_with_scores = [
            (-arb_abs(arb_sub(rank, target_rank, prec)), l, u)
            for rank, l, u in intervals
        ]
        _, l, u = sorted(intervals_with_scores, reverse=True)[0]
        l_float = l.to_float()
        u_float = u.to_float()
        return l_float, u_float

    # need to be calculated w/ arb, calculation on floats can be inexact
    delta_u: Arb = arb_max(arb_q, arb_sub(Arb.from_int(1), arb_q, prec), prec)

    # select bin
    gumbel_p_bits = [0] * len(intervals)
    n = 0

    step_size = 15  # optimal step size from benchmarking
    while len(intervals) > 1:
        n += step_size
        prec = n

        # sample Gumbel noise with more bits
        gumbel_p_bits = [
            (old_bits << step_size) + int(new_bits)
            for old_bits, new_bits in zip(
                gumbel_p_bits,
                prng().integers(pow(2, step_size), size=len(gumbel_p_bits)),
            )
        ]
        probabilities = [
            Arb.from_midpoint_radius(
                mid=Arb.from_man_exp(p_bits, -n), rad=Arb.from_man_exp(1, -n)
            )
            for p_bits in gumbel_p_bits
        ]
        # probabilities for sampling Gumbel noise using the inverse CDF

        gumbels = [-arb_log(-arb_log(p, prec), prec) for p in probabilities]

        noisy_scores = [
            # arb.log(u - l) - ((abs(rank - target_rank) * epsilon) / (2 * delta_u)) + noise
            arb_add(
                arb_sub(
                    arb_log(arb_sub(u, l, prec), prec),
                    arb_div(
                        arb_mul(
                            arb_abs(arb_sub(rank, target_rank, prec)),
                            Arb.from_float(epsilon),
                            prec,
                        ),
                        arb_mul(Arb.from_int(2), delta_u, prec),
                        prec,
                    ),
                    prec,
                ),
                noise,
                prec,
            )
            for noise, (rank, l, u) in zip(gumbels, intervals)
        ]

        # try to get a noisy score which is above most others
        approx_max = Arb.from_float(float("-inf"))
        # pylint: disable=consider-using-max-builtin
        # Unclear if max works correctly with Arb
        for noisy_score in noisy_scores:
            if noisy_score > approx_max:
                # only if noisy_score.lower > approx_max.upper
                approx_max = noisy_score
        # pylint: enable=consider-using-max-builtin

        # do another pass to eliminate other intervals
        new_gumbel_p_bits = []
        remaining_intervals: List[_RankedInterval] = []
        for i, noisy_score in enumerate(noisy_scores):
            if not (
                noisy_score
                < approx_max
                # NOT the same as noisy_score >= approx_max
                # A < B only returns true if A.upper < B.lower
                # true if A.upper < B.lower
            ):
                new_gumbel_p_bits.append(gumbel_p_bits[i])
                remaining_intervals.append(intervals[i])
        gumbel_p_bits = new_gumbel_p_bits
        intervals = remaining_intervals
    assert len(intervals) == 1
    _, l, u = intervals[0]
    return l.to_float(), u.to_float()
