"""RDP computations for different subsampling strategies"""
# pylint: disable=cyclic-import

from typing import TYPE_CHECKING, cast

import numpy as np
from scipy import special

if TYPE_CHECKING:
    from sarus_differential_privacy.query import PrivateQuery


# Subsampling without replacement
def rdp_without_replacement(
    sampling_ratio: float, alpha: float, query: 'PrivateQuery'
) -> float:
    """RDP for subsampling without replacement"""
    if float(alpha).is_integer():
        return wor_compute_eps_int(sampling_ratio, int(alpha), query)
    return wor_compute_eps_frac(sampling_ratio, alpha, query)


def wor_compute_eps_int(
    sampling_ratio: float, alpha: int, query: 'PrivateQuery'
) -> float:
    """Theorem 9 from https://arxiv.org/pdf/1808.00087.pdf"""
    if alpha == 1:
        return 0

    log_inner_log = [0]
    rdp2 = cast(float, query.rdp(2))
    log_inner_log.append(
        2 * np.log(sampling_ratio)
        + np.log(special.binom(alpha, 2))
        + min(
            2 * np.log(2)
            + np.log(special.expm1(rdp2)),  # pylint: disable=no-member
            rdp2
            + min(
                np.log(2),
                2
                * np.log(
                    special.expm1(  # pylint: disable=no-member
                        cast(float, query.rdp(np.inf))
                    )
                ),
            ),
        )
    )

    if alpha >= 3:
        for j in range(3, alpha + 1):
            log_inner_log.append(
                j * np.log(sampling_ratio)
                + np.log(special.binom(alpha, j))
                + (j - 1) * cast(float, query.rdp(j))
                + min(
                    np.log(2),
                    j
                    * np.log(
                        special.expm1(  # pylint: disable=no-member
                            cast(float, query.rdp(np.inf))
                        )
                    ),
                )
            )

    new_eps = (1 / (alpha - 1)) * cast(float, special.logsumexp(log_inner_log))
    return new_eps  # typing: ignore[no-any-return]


def wor_compute_eps_frac(
    sampling_ratio: float, alpha: float, query: 'PrivateQuery'
) -> float:
    """Corollary 10 from https://arxiv.org/pdf/1808.00087.pdf"""
    int_alpha = int(alpha)
    new_eps = (1 - alpha + int_alpha) * (int_alpha - 1) / (
        alpha - 1
    ) * wor_compute_eps_int(sampling_ratio, int_alpha, query) + (
        alpha - int_alpha
    ) * int_alpha / (
        alpha - 1
    ) * wor_compute_eps_int(
        sampling_ratio, int_alpha + 1, query
    )
    return new_eps


# poisson subsampling
def rdp_poisson(
    sampling_ratio: float, alpha: float, query: 'PrivateQuery'
) -> float:
    """RDP for Poisson subsampling"""
    if float(alpha).is_integer():
        return poisson_compute_eps_int(sampling_ratio, int(alpha), query)
    return poisson_compute_eps_frac(sampling_ratio, alpha, query)


def poisson_compute_eps_int(
    sampling_ratio: float, alpha: int, query: 'PrivateQuery'
) -> float:
    """Theorem 5 from http://proceedings.mlr.press/v97/zhu19c/zhu19c.pdf"""
    log_inner_log = [
        (alpha - 1) * np.log(1 - sampling_ratio)
        + np.log((alpha * sampling_ratio) - sampling_ratio + 1)
    ]
    rdp2 = cast(float, query.rdp(2))
    log_inner_log.append(
        2 * np.log(sampling_ratio)
        + np.log(special.binom(alpha, 2))
        + rdp2
        + np.log(1 - sampling_ratio) * (alpha - 2)
    )

    if alpha >= 3:
        for lambd in range(3, alpha + 1):
            log_inner_log.append(
                np.log(3)
                + np.log(special.binom(alpha, lambd))
                + np.log(1 - sampling_ratio) * (alpha - lambd)
                + np.log(sampling_ratio) * lambd
                + (lambd - 1) * cast(float, query.rdp(lambd))
            )

    new_eps = (1 / (alpha - 1)) * cast(float, special.logsumexp(log_inner_log))
    return new_eps  # typing: ignore[no-any-return]


def poisson_compute_eps_frac(
    sampling_ratio: float, alpha: float, query: 'PrivateQuery'
) -> float:
    """Corollary 10 from https://arxiv.org/pdf/1808.00087.pdf"""
    int_alpha = int(alpha)
    new_eps = (1 - alpha + int_alpha) * (int_alpha - 1) / (
        alpha - 1
    ) * wor_compute_eps_int(sampling_ratio, int_alpha, query) + (
        alpha - int_alpha
    ) * int_alpha / (
        alpha - 1
    ) * wor_compute_eps_int(
        sampling_ratio, int_alpha + 1, query
    )
    return new_eps