"""Privacy risk metrics and conversions.

This module provides functions for computing privacy metrics (beta, advantage,
Bayes risk) from different privacy representations (PLDs, ADP, GDP,
RDP, zCDP).
"""

from typing import Union, Optional
import warnings
import numpy as np
from scipy import stats, optimize
from dp_accounting.pld import privacy_loss_distribution

from riskcal.analysis import _plrv
from riskcal.utils import _ensure_array

from riskcal.analysis.rdp import get_FNR as _get_FNR_from_rdp


# =============================================================================
# Internal conversions
# =============================================================================


def pld_to_plrvs(
    pld: privacy_loss_distribution.PrivacyLossDistribution,
) -> _plrv.PLRVs:
    """
    Convert PLD to internal PLRVs representation.

    Converts a Google dp_accounting Privacy Loss Distribution object
    into the PLRVs (Privacy Loss Random Variables) format used internally.

    Args:
        pld: Privacy loss distribution from Google's dp_accounting library.

    Returns:
        PLRVs object containing the privacy loss random variables.

    Note:
        This is an internal conversion function. Most users should use
        the higher-level `get_beta_from_pld()`, `get_advantage_from_pld()`, etc.
    """

    def _get_plrv(pld):
        pld = pld.to_dense_pmf()
        pmf = pld._probs
        lower_loss = pld._lower_loss
        infinity_mass = pld._infinity_mass
        return lower_loss, infinity_mass, pmf

    lower_loss_Y, infinity_mass_Y, pmf_Y = _get_plrv(pld._pmf_remove)
    lower_loss_Z, infinity_mass_Z, pmf_Z = _get_plrv(pld._pmf_add)

    upper_loss_Z = lower_loss_Z + len(pmf_Z) - 1

    # clean pmfs. Sometimes float errors cause probs to be negative
    pmf_Y = np.where(pmf_Y < 0, 0, pmf_Y)
    pmf_Z = np.where(pmf_Z < 0, 0, pmf_Z)

    pmf_Y = pmf_Y * (1 - infinity_mass_Y) / np.sum(pmf_Y)
    pmf_Z = pmf_Z * (1 - infinity_mass_Z) / np.sum(pmf_Z)

    is_symmetric = pld._symmetric
    return _plrv.PLRVs(
        y0=lower_loss_Y,
        x0=-upper_loss_Z,
        pmf_Y=pmf_Y,
        pmf_X=pmf_Z[::-1],
        minus_infinity_mass_X=infinity_mass_Z,
        infinity_mass_Y=infinity_mass_Y,
        is_symmetric=is_symmetric,
    )


# Backward compatibility alias
plrvs_from_pld = pld_to_plrvs
plrvs_from_pld.__deprecated__ = "2.0.0"


# =============================================================================
# PLD (Privacy Loss Distribution)
# =============================================================================


def get_beta_from_pld(
    pld: privacy_loss_distribution.PrivacyLossDistribution,
    alpha: Union[float, np.ndarray] = None,
    alphas: Union[float, np.ndarray] = None,  # Deprecated
) -> Union[float, np.ndarray]:
    """
    Compute false negative rate (FNR) for given false positive rate (FPR) from PLD.

    Uses the direct method from Algorithm 1 (Kulynych et al., 2024) to compute
    the optimal trade-off between FNR (beta) and FPR (alpha).

    .. deprecated:: 1.2.0
        Parameter 'alphas' is deprecated and will be removed in version 2.0.0.
        Use 'alpha' instead.

    Args:
        pld: Privacy loss distribution from Google's dp_accounting library.
        alpha: False positive rate(s) in [0, 1]. Can be scalar or array.
        alphas: (Deprecated) Use 'alpha' instead.

    Returns:
        False negative rate(s) corresponding to input alpha.

    Example:
        >>> from dp_accounting.pld import privacy_loss_distribution as pld_module
        >>> pld = pld_module.from_gaussian_mechanism(1.0)
        >>> beta = get_beta_from_pld(pld, alpha=0.01)

    References:
        Kulynych & Gomez et al. (2024), Algorithm 1. https://arxiv.org/abs/2407.02191
    """
    if alpha is None and alphas is None:
        raise ValueError("Must specify alpha.")
    elif alpha is not None and alphas is not None:
        raise ValueError("Must pass either alpha or alphas.")
    elif alphas is not None:
        warnings.warn(
            "Parameter 'alphas' is deprecated. Use 'alpha' instead.",
            DeprecationWarning,
            stacklevel=2,
        )
        alpha = alphas

    return _plrv.get_beta(pld_to_plrvs(pld), alpha)


def get_advantage_from_pld(
    pld: privacy_loss_distribution.PrivacyLossDistribution,
) -> float:
    """
    Compute attack advantage from PLD.

    Advantage is the maximum value of (TPR - FPR) achievable by any attacker,
    which equals delta at epsilon=0.

    Args:
        pld: Privacy loss distribution from Google's dp_accounting library.

    Returns:
        Maximum attack advantage.

    References:
        Kulynych & Gomez et al. (2024). https://arxiv.org/abs/2407.02191
    """
    return pld.get_delta_for_epsilon(0)


def get_bayes_risk_from_pld(pld, prior):
    """
    Compute Bayes Risk from PLD for given prior probability.

    Bayes risk shows the maximum accuracy of an attack against privacy of a single
    record under a binary prior (e.g., accuracy of attribute inference).

    Args:
        pld: Privacy loss distribution from Google's dp_accounting library.
        prior: Prior probability (scalar or array). Probability that the
            sensitive attribute takes value 1.

    Returns:
        Bayes risk value(s). Float if prior is scalar, array if prior is array.

    Example:
        >>> from dp_accounting.pld import privacy_loss_distribution as pld_module
        >>> pld = pld_module.from_laplace_mechanism(1.0)
        >>> risk = get_bayes_risk_from_pld(pld, prior=0.5)

    References:
        Kulynych et al. (2025), Proposition D.1 / Eq. 35.
        https://arxiv.org/abs/2507.06969
    """
    prior, is_scalar = _ensure_array(prior)
    epsilon = np.log(prior / (1 - prior))
    delta = np.array([pld.get_delta_for_epsilon(e) for e in epsilon])
    bayes_risk = (1 - prior) * (1 - delta)

    if is_scalar:
        return bayes_risk[0]
    return np.array(bayes_risk)


# =============================================================================
# GDP (Gaussian Differential Privacy)
# =============================================================================


def get_beta_from_gdp(
    mu: float, alpha: Union[float, np.ndarray]
) -> Union[float, np.ndarray]:
    """
    Compute FNR for FPR using the analytical formula for Gaussian DP.

    Args:
        mu: Gaussian noise scale parameter (sigma).
        alpha: False positive rate(s) in [0, 1].

    Returns:
        False negative rate(s) corresponding to input alpha.

    References:
        Dong et al. (2019), Eq. 6. https://arxiv.org/abs/1905.02383
    """
    return stats.norm.cdf(-stats.norm.ppf(alpha) - mu)


def get_advantage_from_gdp(mu: Union[float, np.ndarray]) -> Union[float, np.ndarray]:
    """
    Compute attack advantage using analytical formula for Gaussian mechanism.

    Args:
        mu: Gaussian noise scale parameter (sigma). Can be scalar or array.

    Returns:
        Attack advantage value(s).

    References:
        Dong et al. (2019), Corollary 2.13. https://arxiv.org/abs/1905.02383
    """
    return stats.norm.cdf(mu / 2) - stats.norm.cdf(-mu / 2)


def get_bayes_risk_from_gdp(mu, prior):
    """
    Compute Bayes Risk for Gaussian mechanism using analytical formula.

    Args:
        mu: Gaussian noise scale parameter (sigma).
        prior: Prior probability (scalar or array).

    Returns:
        Bayes risk value(s). Float if prior is scalar, array if prior is array.

    References:
        Kulynych et al. (2025), Proposition D.1 / Eq. 35.
        https://arxiv.org/abs/2507.06969
        Dong et al. (2019), Corollary 2.13. https://arxiv.org/abs/1905.02383
    """
    assert mu >= 0, "mu must be >= 0"
    prior, is_scalar = _ensure_array(prior)
    epsilon = np.log(prior / (1 - prior))
    # Gaussian privacy profile: Dong et al. (2019), Corollary 2.13
    delta = stats.norm.cdf(-epsilon / mu + mu / 2) - np.exp(epsilon) * stats.norm.cdf(
        -epsilon / mu - mu / 2
    )
    bayes_risk = (1 - prior) * (1 - delta)

    if is_scalar:
        return bayes_risk[0]
    return np.array(bayes_risk)


# Backward compatibility aliases
get_beta_for_mu = get_beta_from_gdp
get_beta_for_mu.__deprecated__ = "2.0.0"
get_advantage_for_mu = get_advantage_from_gdp
get_advantage_for_mu.__deprecated__ = "2.0.0"
get_bayes_risk_for_mu = get_bayes_risk_from_gdp
get_bayes_risk_for_mu.__deprecated__ = "2.0.0"


# =============================================================================
# ADP (Approximate Differential Privacy)
# =============================================================================


def get_beta_from_adp(
    epsilon: float, delta: float, alpha: Union[float, np.ndarray]
) -> Union[float, np.ndarray]:
    """
    Compute FNR for FPR from (epsilon, delta)-DP parameters.

    Args:
        epsilon: Privacy parameter epsilon.
        delta: Privacy parameter delta.
        alpha: False positive rate(s) in [0, 1].

    Returns:
        False negative rate(s) corresponding to input alpha.

    Example:
        >>> import numpy as np
        >>> np.round(get_beta_from_adp(1.0, 0.001, 0.8), 3)
        0.073

    References:
        Dong et al. (2019), Eq. 5. https://arxiv.org/abs/1905.02383
    """
    form1 = np.array(1 - delta - np.exp(epsilon) * alpha)
    form2 = np.array(np.exp(-epsilon) * (1 - delta - alpha))
    return np.maximum.reduce([form1, form2, np.zeros_like(form1)])


def get_advantage_from_adp(epsilon: float, delta: float) -> float:
    """
    Compute attack advantage from (epsilon, delta)-DP parameters.

    Args:
        epsilon: Privacy parameter epsilon.
        delta: Privacy parameter delta.

    Returns:
        Attack advantage.

    Example:
        >>> import numpy as np
        >>> np.round(get_advantage_from_adp(0., 0.001), 3)
        0.001

    References:
        Dong et al. (2019). https://arxiv.org/abs/1905.02383
    """
    return (np.exp(epsilon) + 2 * delta - 1) / (np.exp(epsilon) + 1)


def get_epsilon_from_err_rates(delta: float, alpha: float, beta: float) -> float:
    """
    Convert f-DP error rates (alpha, beta) to epsilon for (epsilon, delta)-DP.

    Args:
        delta: Target delta parameter for (epsilon, delta)-DP.
        alpha: False positive rate (FPR) from f-DP.
        beta: False negative rate (FNR) from f-DP.

    Returns:
        Epsilon value corresponding to the error rates.

    Example:
        >>> import numpy as np
        >>> np.round(get_epsilon_from_err_rates(0.001, 0.001, 0.8), 3)
        5.293

    References:
        Dong et al. (2019). https://arxiv.org/abs/1905.02383
    """
    epsilon1 = np.log((1 - delta - alpha) / beta)
    epsilon2 = np.log((1 - delta - beta) / alpha)
    return np.maximum.reduce([epsilon1, epsilon2, np.zeros_like(epsilon1)])


# Backward compatibility aliases
get_beta_for_epsilon_delta = get_beta_from_adp
get_beta_for_epsilon_delta.__deprecated__ = "2.0.0"
get_advantage_for_epsilon_delta = get_advantage_from_adp
get_advantage_for_epsilon_delta.__deprecated__ = "2.0.0"
get_epsilon_for_err_rates = get_epsilon_from_err_rates
get_epsilon_for_err_rates.__deprecated__ = "2.0.0"


# =============================================================================
# RDP (Renyi Differential Privacy)
# =============================================================================


def get_beta_from_rdp(
    epsilon: float,
    alpha: Union[float, np.ndarray],
    order: float,
    linear_search_step: float = 1e-3,  # unused; kept for backward compatibility
    max_bisection_steps: int = 50,  # unused; kept for backward compatibility
    tol: float = 1e-7,
) -> Union[float, np.ndarray]:
    """
    Compute FNR for FPR from Renyi DP parameters.

    Uses the optimal conversion from a single RDP guarantee to a tradeoff
    function (Riess et al., 2026). Accepts scalar or array FPR values.

    Args:
        epsilon: Renyi divergence parameter (privacy budget).
        alpha: False positive rate(s) (FPR) in [0, 1]. Scalar or array.
        order: Order of Renyi divergence (alpha in Renyi DP literature).
        linear_search_step: Unused. Kept for backward compatibility.
        max_bisection_steps: Unused. Kept for backward compatibility.
        tol: Numerical tolerance for bisection convergence.

    Returns:
        False negative rate (FNR) corresponding to input alpha.

    References:
        Zhu et al. (2022), Appendix F.1. https://arxiv.org/abs/2106.08567
        Riess et al. (2026). https://arxiv.org/abs/2602.04562
    """
    return _get_FNR_from_rdp(alpha_array=alpha, order=order, epsilon=epsilon, tol=tol)


# =============================================================================
# zCDP (Zero-Concentrated Differential Privacy)
# =============================================================================


def get_beta_from_zcdp(
    rho: float,
    alpha: Union[float, np.ndarray],
    tol: float = 1e-9,
    min_order: float = 1.0,
    max_order: float = 1024.0,
    grid_size=1000,
) -> Union[float, np.ndarray]:
    """
    Compute FNR for FPR from zCDP parameter rho.

    Zero-Concentrated Differential Privacy (zCDP) is characterized by a single
    parameter rho. This function computes the optimal trade-off between false
    negative rate (beta) and false positive rate (alpha) by optimizing over
    Renyi divergence orders.

    Args:
        rho: Zero-concentrated differential privacy parameter.
        alpha: False positive rate(s) in [0, 1]. Can be scalar or array.
        tol: Tolerance for numerical convergence and avoiding log(0).
        max_order: Maximum Renyi DP order to consider.

    Returns:
        False negative rate(s) corresponding to input alpha. Returns float if
        alpha is scalar, array if alpha is array.

    Example:
        >>> import numpy as np
        >>> # Single alpha value
        >>> beta = get_beta_from_zcdp(rho=0.5, alpha=0.1)
        >>> np.round(beta, 3)
        0.517
        >>> # Multiple alpha values
        >>> betas = get_beta_from_zcdp(rho=0.5, alpha=np.array([0.1, 0.2, 0.3]))
        >>> np.round(betas, 3)
        array([0.517, 0.34, 0.232])

    References:
        Bun & Steinke (2016). https://arxiv.org/abs/1605.02065
        Zhu et al. (2022), Appendix F.1. https://arxiv.org/abs/2106.08567
        Riess et al. (2026). https://arxiv.org/abs/2602.04562
    """
    scalar_input = isinstance(alpha, (int, float))
    alpha_arr = np.atleast_1d(np.asarray(alpha, dtype=float))

    result = np.empty_like(alpha_arr)
    low = alpha_arr <= tol
    high = alpha_arr >= 1 - tol
    interior = ~low & ~high

    result[low] = 1.0
    result[high] = 0.0

    orders_grid = np.concatenate(
        [
            np.linspace(min_order, 2, grid_size // 3)[:-1],
            np.linspace(2, 16, grid_size // 3)[:-1],
            np.logspace(np.log10(16), np.log10(max_order), grid_size // 3),
        ]
    )
    if interior.any():
        alpha_interior = alpha_arr[interior]
        best_vals = np.zeros(alpha_interior.shape)
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            for order in orders_grid:
                best_vals = np.maximum(
                    best_vals,
                    np.atleast_1d(
                        get_beta_from_rdp(
                            epsilon=rho * order,
                            alpha=alpha_interior,
                            order=order,
                            tol=tol,
                        )
                    ),
                )
        result[interior] = best_vals

    return float(result[0]) if scalar_input else result


def get_advantage_from_zcdp(
    rho: float,
    tol: float = 1e-9,
    min_order: float = 1.0,
    max_order: float = 1024.0,
    grid_size: int = 1000,
) -> float:
    """
    Compute attack advantage from zCDP parameter rho.

    Advantage is the maximum value of (TPR - FPR) achievable by any attacker.
    This function optimizes over all possible threshold choices to find the
    maximum advantage.

    Args:
        rho: Zero-concentrated differential privacy parameter.
        tol: Tolerance for numerical convergence and avoiding log(0).
        min_order: Minimum Renyi DP order to consider.
        max_order: Maximum Renyi DP order to consider.
        grid_size: Number of orders to evaluate in the grid search.

    Returns:
        Maximum attack advantage.

    Example:
        >>> import numpy as np
        >>> adv = get_advantage_from_zcdp(rho=0.5)
        >>> np.round(adv, 3)
        0.47

    References:
        Bun & Steinke (2016). https://arxiv.org/abs/1605.02065
    """
    result = optimize.minimize_scalar(
        lambda alpha: -(
            1 - alpha - get_beta_from_zcdp(
                rho=rho, alpha=alpha, tol=tol, min_order=min_order,
                max_order=max_order, grid_size=grid_size,
            )
        ),
        bounds=(0, 1),
        method="bounded",
        options={"xatol": 1e-12},
    )
    if not result.success:
        warnings.warn("Optimization failed for advantage calculation")
        return np.nan
    else:
        return -result.fun


def get_mu_from_zcdp_approx(rho: float) -> float:
    """Deprecated. Use get_beta_from_zcdp or get_advantage_from_zcdp instead."""
    warnings.warn(
        "get_mu_from_zcdp_approx is deprecated and will be removed in a future version. "
        "Use get_beta_from_zcdp or get_advantage_from_zcdp for exact computations.",
        DeprecationWarning,
        stacklevel=2,
    )
    return (
        1.5822558654881096 * np.sqrt(rho)
        + 0.08064620797681155 * rho
        + 0.05485600531538526
    )


def get_beta_from_zcdp_approx(
    rho: float, alpha: Union[float, np.ndarray]
) -> Union[float, np.ndarray]:
    """Deprecated. Use get_beta_from_zcdp instead."""
    warnings.warn(
        "get_beta_from_zcdp_approx is deprecated and will be removed in a future version. "
        "Use get_beta_from_zcdp for exact computation.",
        DeprecationWarning,
        stacklevel=2,
    )
    return get_beta_from_zcdp(rho, alpha=alpha)


def get_advantage_from_zcdp_approx(rho: float) -> float:
    """Deprecated. Use get_advantage_from_zcdp instead."""
    warnings.warn(
        "get_advantage_from_zcdp_approx is deprecated and will be removed in a future version. "
        "Use get_advantage_from_zcdp for exact computation.",
        DeprecationWarning,
        stacklevel=2,
    )
    return get_advantage_from_zcdp(rho)
