"""Adaptive helpers for random-allocation queries."""

from __future__ import annotations

import math
import warnings
from dataclasses import dataclass
from typing import Callable, Optional

import numpy as np

from PLD_accounting.types import AllocationSchemeConfig, BoundType, ConvolutionMethod, Direction, PrivacyParams
from dp_accounting.pld import privacy_loss_distribution


# Global configuration constants
MAX_ITERATIONS = 10
POISSON_GUESS_DISCRETIZATION = 1e-4
MIN_ESTIMATE = 1e-300
MIN_DISCRETIZATION = float(np.finfo(float).eps)
MIN_TAIL_TRUNCATION = float(np.finfo(float).tiny)


@dataclass
class AdaptiveResult:
    """Result from adaptive allocation computation.

    Attributes:
        upper_bound: Best upper bound found across all iterations.
        lower_bound: Best lower bound found across all iterations.
        absolute_gap: Final gap between upper and lower bounds.
        converged: Whether the algorithm converged to target accuracy.
        iterations: Number of refinement iterations performed.
        initial_discretization: Starting loss_discretization value.
        discretization: Final loss_discretization value used.
        initial_tail_truncation: Starting tail_truncation value.
        tail_truncation: Final tail_truncation value used.
        target_accuracy: Target absolute gap for convergence.
    """
    upper_bound: float
    lower_bound: float
    absolute_gap: float
    converged: bool
    iterations: int
    initial_discretization: float
    discretization: float
    initial_tail_truncation: float
    tail_truncation: float
    target_accuracy: float


def estimate_poisson_query(
    params: PrivacyParams,
    query_func: Callable[[privacy_loss_distribution.PrivacyLossDistribution], float],
) -> float:
    """Estimate the query value with a Poisson-subsampled Gaussian approximation."""
    sampling_probability = params.num_selected / params.num_steps
    compose_steps = params.num_selected * params.num_epochs

    pld = privacy_loss_distribution.from_gaussian_mechanism(
        standard_deviation=params.sigma,
        sensitivity=1.0,
        value_discretization_interval=POISSON_GUESS_DISCRETIZATION,
        pessimistic_estimate=True,
        sampling_prob=sampling_probability,
    ).self_compose(compose_steps)

    estimate = float(query_func(pld))
    if not math.isfinite(estimate) or estimate < 0.0:
        raise RuntimeError(
            "Poisson-based adaptive initialization produced an invalid estimate: "
            f"{estimate!r}"
        )
    return max(estimate, MIN_ESTIMATE)


def adaptive_convergence(
    params: PrivacyParams,
    target_accuracy: float,
    initial_discretization: float,
    initial_tail_truncation: float,
    query_func: Callable[[privacy_loss_distribution.PrivacyLossDistribution], float],
    tail_effect_func: Callable[[privacy_loss_distribution.PrivacyLossDistribution, float], float],
    pld_builder: Callable[..., privacy_loss_distribution.PrivacyLossDistribution],
) -> AdaptiveResult:
    """Core adaptive refinement loop.

    Args:
        params: Privacy parameters
        target_accuracy: Target absolute gap between bounds
        query_func: Function to extract the scalar query value from the computed PLD.
        tail_effect_func: Estimates how much of the current query gap could come
            from the current tail truncation budget using the upper PLD.
        pld_builder: Callable that constructs the PLD for a given config/bound pair.

    Returns:
        AdaptiveResult with convergence metadata
    """
    discretization = max(initial_discretization, MIN_DISCRETIZATION)
    tail_truncation = max(initial_tail_truncation, MIN_TAIL_TRUNCATION)

    upper_bound = np.inf
    lower_bound = -np.inf
    converged = False
    iteration = 0

    while not converged and iteration < MAX_ITERATIONS:
        iteration += 1
        config = AllocationSchemeConfig(
            loss_discretization=discretization,
            tail_truncation=tail_truncation,
            convolution_method=ConvolutionMethod.GEOM,
        )

        # Compute upper bound (DOMINATES)
        pld_upper = pld_builder(
            params=params,
            config=config,
            direction=Direction.BOTH,
            bound_type=BoundType.DOMINATES,
        )
        new_upper = query_func(pld_upper)

        # Compute lower bound (IS_DOMINATED)
        pld_lower = pld_builder(
            params=params,
            config=config,
            direction=Direction.BOTH,
            bound_type=BoundType.IS_DOMINATED,
        )
        new_lower = query_func(pld_lower)

        if new_upper < new_lower:
            raise RuntimeError(
                "Adaptive refinement produced invalid bounds: dominating bound "
                f"{new_upper:.12g} is below dominated bound {new_lower:.12g}"
            )

        upper_bound = min(upper_bound, new_upper)
        lower_bound = max(lower_bound, new_lower)

        if upper_bound < lower_bound:
            raise RuntimeError(
                "Adaptive refinement produced invalid bounds: dominating bound "
                f"{upper_bound:.12g} is below dominated bound {lower_bound:.12g}"
            )

        # Check convergence
        gap = upper_bound - lower_bound
        if gap < target_accuracy:
            converged = True
            break

        current_gap = new_upper - new_lower
        estimated_tail_effect = float(tail_effect_func(pld_upper, tail_truncation))
        if not math.isfinite(estimated_tail_effect) or estimated_tail_effect < 0.0:
            raise RuntimeError(
                "Adaptive refinement produced an invalid tail-effect estimate: "
                f"{estimated_tail_effect!r}"
            )

        if current_gap > estimated_tail_effect:
            discretization = max(discretization / 2, MIN_DISCRETIZATION)
        else:
            tail_truncation = max(tail_truncation / 10, MIN_TAIL_TRUNCATION)

    # Warn if not converged
    if not converged:
        warnings.warn(
            f"Adaptive refinement did not converge after {MAX_ITERATIONS} iterations. "
            f"Final gap: {upper_bound - lower_bound:.6e}, target: {target_accuracy:.6e}. "
            f"Returning best bounds found.",
            RuntimeWarning,
        )

    return AdaptiveResult(
        upper_bound=upper_bound,
        lower_bound=lower_bound,
        absolute_gap=upper_bound - lower_bound,
        converged=converged,
        iterations=iteration,
        initial_discretization=max(initial_discretization, MIN_DISCRETIZATION),
        discretization=discretization,
        initial_tail_truncation=max(initial_tail_truncation, MIN_TAIL_TRUNCATION),
        tail_truncation=tail_truncation,
        target_accuracy=target_accuracy,
    )
