"""Privacy Accountant classes"""
# pylint: disable=redefined-builtin, too-many-lines

from __future__ import annotations

from typing import Collection, Optional, Sequence, Tuple, Union, cast

import numpy as np

from sarus_differential_privacy.accountant import PrivateQueryFilter
from sarus_differential_privacy.accountant_local import LocalPrivacyAccountant
from sarus_differential_privacy.query import PrivateQuery
from sarus_differential_privacy.typing import PrivacyLimit


# pylint: disable=no-self-use, too-many-public-methods
class LocalPrivacyOdometer(LocalPrivacyAccountant):
    """The privacy odometer from https://arxiv.org/pdf/2103.01379.pdf, based
    on the implementation from https://github.com/m-lecuyer/adaptive_rdp/blob/
    main/adaptive_privacy_engine.py.
    delta() methods are not implemented here because by definition odometers
    cannot have delta() methods.
    Consequently, verify() methods only verify the limits on epsilon and
    pure_delta.
    """

    def __init__(
        self,
        alphas: Optional[Sequence[float]] = None,
    ):
        super().__init__(alphas)
        self.filter_order = 1

    def filter_family(self, filter_order: int, delta: float) -> np.ndarray:
        return cast(
            np.ndarray,
            2 ** (filter_order - 1)
            * np.log(2 * len(self.alphas) / delta)
            / (np.atleast_1d(self.alphas) - 1),
        )

    def _epsilon_from_rdp(self, rdp: Sequence[float], delta: float) -> float:
        gamma = (
            2**-2
            * np.log(2 * len(self.alphas) / delta)
            / (np.atleast_1d(self.alphas) - 1)
        )

        if all(v == 0 for v in rdp):
            return 0

        rdp_max = np.maximum(rdp, gamma)
        f = np.ceil(np.log2(rdp_max / gamma))
        rdp_max = gamma * np.exp2(f)
        return self.get_privacy_spent_heterogeneous_delta(
            self.alphas,
            rdp_max.tolist(),
            delta / (len(self.alphas) * 2 * np.power(f + 1, 2)),
        )[0]

    def get_privacy_spent_heterogeneous_delta(
        self,
        orders: Union[Sequence[float], float],
        rdp: Union[Sequence[float], float],
        delta: Union[Sequence[float], float],
    ) -> Tuple[float, float]:
        """Computes epsilon given a list of Renyi Differential Privacy (RDP)
        values at multiple RDP orders and target ``delta``.
        Args:
            orders: An array (or a scalar) of orders (alphas).
            rdp: A list (or a scalar) of RDP guarantees.
            delta: A list (or a scalar) of target deltas for each order.
        Returns:
            Pair of epsilon and optimal order alpha.
        Raises:
            ValueError
                If the lengths of ``orders`` and ``rdp`` are not equal.
        """
        orders_vec = np.atleast_1d(orders)
        rdp_vec = np.atleast_1d(rdp)
        delta_vec = np.atleast_1d(delta)

        if len(orders_vec) != len(rdp_vec) or len(orders_vec) != len(
            delta_vec
        ):
            raise ValueError(
                f"Input lists must have the same length.\n"
                f"\torders_vec = {len(orders_vec)}\n"
                f"\trdp_vec = {len(rdp_vec)}\n"
                f"\tdelta_vec = {len(delta_vec)}\n"
            )

        eps = rdp_vec - np.log(delta) / (orders_vec - 1)

        # special case when there is no privacy
        if np.isnan(eps).all():
            return np.inf, np.nan

        idx_opt = np.nanargmin(eps)  # Ignore NaNs
        return eps[idx_opt], orders_vec[idx_opt]

    def delta(
        self, epsilon: float, filter: Optional[PrivateQueryFilter] = None
    ) -> float:
        raise NotImplementedError

    def delta_new_query(
        self,
        epsilon: float,
        query: Union[PrivateQuery, Collection[PrivateQuery]],
        filter: Optional[PrivateQueryFilter] = None,
    ) -> float:
        """Delta of accountant + query 'query'"""
        raise NotImplementedError

    def delta_query(
        self,
        epsilon: float,
        query: Union[PrivateQuery, Collection[PrivateQuery]],
    ) -> float:
        """Delta of single query 'query'"""
        raise NotImplementedError

    def verify(
        self,
        limit: Optional[PrivacyLimit],
        filter: Optional[PrivateQueryFilter] = None,
    ) -> bool:
        """Verify if the (filtered) ledger goes over the limit"""
        if not limit:
            return True

        for delta, epsilon in limit.delta_epsilon_dict().items():
            if (
                self.epsilon(delta, filter) > epsilon
                or self.pure_delta(filter) > delta
            ):
                return False
        return True

    def verify_new_query(
        self,
        query: Union[PrivateQuery, Collection[PrivateQuery]],
        limit: Optional[PrivacyLimit],
        filter: Optional[PrivateQueryFilter] = None,
    ) -> bool:
        """Verify if a query can be added to the ledger without actually
        adding it"""
        if not limit:
            return True

        for delta, epsilon in limit.delta_epsilon_dict().items():

            if (
                self.epsilon_new_query(delta, query, filter) > epsilon
                or self.pure_delta_new_query(query, filter) > delta
            ):
                return False
        return True

    def verify_query(
        self,
        query: Union[PrivateQuery, Collection[PrivateQuery]],
        limit: Optional[PrivacyLimit],
    ) -> bool:
        """Verify if a query (without the ledger) passes the limit"""
        if not limit:
            return True

        for delta, epsilon in limit.delta_epsilon_dict().items():
            if (
                self.epsilon_query(delta, query) > epsilon
                or self.pure_delta_query(query) > delta
            ):
                return False
        return True