import math
import random
from typing import Iterable, List, Optional, Tuple

############################
# Utility: noise samplers  #
############################

def _gumbel(scale: float, rng: random.Random) -> float:
    # Gumbel(0, scale) via inverse CDF: -scale * log(-log U)
    # add tiny eps for numerical safety
    u = max(1e-12, min(1 - 1e-12, rng.random()))
    return -scale * math.log(-math.log(u))

def _laplace(scale: float, rng: random.Random) -> float:
    # Laplace(0, scale) via inverse CDF from U(-0.5, 0.5)
    u = rng.random() - 0.5
    return -scale * math.copysign(1.0, u) * math.log(1 - 2 * abs(u))

###############################################
# (A) One-shot Gumbel Top-k (Exponential Mech)#
###############################################

def dp_topk_gumbel_one_shot(
    scores: Iterable[float],
    k: int,
    epsilon: float,
    sensitivity: float = 1.0,
    *,
    bounded_range: bool = True,
    seed: Optional[int] = None,
    return_noisy: bool = False,
    mode: str = "top",          
) -> Tuple[List[int], Optional[List[Tuple[int, float]]]]:

    if mode not in ("top", "bottom"):
        raise ValueError("mode must be 'top' or 'bottom'")

    rng = random.Random(seed)
    s = list(scores)
    n = len(s)
    k = min(k, n)
    if k <= 0 or n == 0:
        return ([], [] if return_noisy else None)

    beta = (sensitivity / epsilon) if bounded_range else (2.0 * sensitivity / epsilon)

    noisy = [(i, s[i] + _gumbel(beta, rng)) for i in range(n)]
    
    noisy_sorted = sorted(noisy, key=lambda t: t[1], reverse=(mode == "top"))
    chosen = noisy_sorted[:k]
    indices = [i for i, _ in chosen]

    return (indices, noisy if return_noisy else None)

#################################################
# (B) One-shot Laplace Top-k (approx (ε,δ)-DP)  #
#################################################

def dp_topk_laplace_one_shot(
    scores: Iterable[float],
    k: int,
    epsilon: float,
    delta: float,
    sensitivity: float = 1.0,
    *,
    seed: Optional[int] = None,
    return_noisy: bool = False,
    release_noisy_values: bool = False,
) -> Tuple[List[int], Optional[List[Tuple[int, float]]], Optional[List[Tuple[int, float]]]]:

    rng = random.Random(seed)
    s = list(scores)
    n = len(s)
    k = min(k, n)
    if k <= 0 or n == 0:
        return [], [] if return_noisy else None, [] if release_noisy_values else None
    if not (0 < delta < 1):
        raise ValueError("delta must be in (0,1)")
    if n <= 1:

        delta = max(delta, 1e-12)

    lam = 8.0 * sensitivity * math.sqrt(k * math.log(n / delta)) / epsilon

    noisy = [(i, s[i] + _laplace(lam, rng)) for i in range(n)]
    noisy.sort(key=lambda t: t[1], reverse=True)
    topk = noisy[:k]
    indices = [i for i, _ in topk]

    noisy_topk_values = None
    if release_noisy_values:
        noisy_topk_values = [(i, s[i] + _laplace(lam, rng)) for i in indices]

    return (indices,
            noisy if return_noisy else None,
            noisy_topk_values)


def dp_topk_gumbel_peeling(
    scores: Iterable[float],
    k: int,
    epsilon_total: float,
    sensitivity: float = 1.0,
    *,
    bounded_range: bool = True,
    seed: Optional[int] = None,
) -> List[int]:

    rng = random.Random(seed)
    s = list(scores)
    n = len(s)
    k = min(k, n)
    if k <= 0 or n == 0:
        return []

    eps_step = epsilon_total / k
    beta = (sensitivity / eps_step) if bounded_range else (2.0 * sensitivity / eps_step)

    remaining = list(range(n))
    chosen: List[int] = []
    for _ in range(k):
        noisy = [(i, s[i] + _gumbel(beta, rng)) for i in remaining]
        noisy.sort(key=lambda t: t[1], reverse=True)
        pick = noisy[0][0]
        chosen.append(pick)
        remaining.remove(pick)
    return chosen