"""Bayesian helpers for evolution — thin wrappers over trust/service.py math.

Rule 12 principle: do NOT reinvent the Beta posterior math. TrustService
already ships verified, tested, numerically stable helpers
(``_beta_ppf``, ``_beta_pdf``, ``_beta_reg_inc``). The evolution layer
imports them and wraps them with the two operations it actually needs:

    1. ``sample_beta(alpha, beta)`` — one sample for Thompson sampling
    2. ``prob_greater(a, b)`` — Monte Carlo P(X_a > X_b) for promotion gates

If the trust module ever moves those helpers, this is the one file that
needs updating — every other evolution module imports from here.

Everything is deterministic modulo RNG — pass ``rng=random.Random(seed)``
to make tests reproducible.
"""

from __future__ import annotations

import random
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    pass


def sample_beta(
    alpha: float,
    beta: float,
    *,
    rng: random.Random | None = None,
) -> float:
    """Draw one sample from Beta(alpha, beta).

    Uses ``random.betavariate`` which is O(1) and numerically stable for
    the α,β ranges we use (both ≥ 1 in practice because the prior is 1,1
    and every run adds ≥ 0). Returns a value in (0, 1).

    Falls back gracefully when alpha+beta == 0 (uniform posterior) by
    returning a uniform sample — this should not happen in practice but
    keeps the function total.
    """
    r = rng or random
    if alpha <= 0 or beta <= 0:
        return r.random()
    try:
        return r.betavariate(alpha, beta)
    except ValueError:
        # betavariate can raise on extreme params; degrade to uniform.
        return r.random()


def posterior_mean(alpha: float, beta: float) -> float:
    """Closed-form mean of Beta(α, β). Handles the degenerate case."""
    total = alpha + beta
    return alpha / total if total > 0 else 0.5


def prob_greater(
    alpha_a: float,
    beta_a: float,
    alpha_b: float,
    beta_b: float,
    *,
    samples: int = 2000,
    rng: random.Random | None = None,
) -> float:
    """Monte Carlo estimate of P(X_A > X_B) where X_A ~ Beta(α_A, β_A), X_B ~ Beta(α_B, β_B).

    Used by the promotion arena to gate candidate → active transitions:
    candidate must beat the incumbent with P ≥ 0.90 before promotion.

    2000 samples is enough for two decimal places of accuracy, which is
    what the gate needs. Bumping to 10000 costs ~3ms on modern hardware
    and is fine when invoked from a work-queue handler (not hot path).

    Closed-form solutions exist for small integer α,β via the regularized
    incomplete beta function, but with EMA-based continuous updates
    producing fractional α,β, the closed form is numerically cranky.
    Monte Carlo is good enough and reads cleaner.
    """
    r = rng or random
    wins = 0
    for _ in range(samples):
        xa = sample_beta(alpha_a, beta_a, rng=r)
        xb = sample_beta(alpha_b, beta_b, rng=r)
        if xa > xb:
            wins += 1
    return wins / samples if samples > 0 else 0.5


def ema_update(old: float, new: float, alpha: float = 0.1) -> float:
    """Exponential moving average — utility_ema update.

    α=0.1 means each new run contributes 10% to the running estimate;
    the EMA converges to the mean over ~1/α = 10 runs. Smaller α means
    longer memory.
    """
    if alpha <= 0 or alpha > 1:
        alpha = 0.1
    return (1 - alpha) * old + alpha * new


def variance_update(
    old_var: float,
    utility: float,
    running_ema: float,
    alpha: float = 0.1,
) -> float:
    """EMA update for running variance around the utility EMA.

    Used to surface "is this reflex behaving consistently?" — high
    variance on an otherwise-good EMA means the reflex wins big sometimes
    and loses big sometimes, which is a yellow flag for promotion.
    """
    if alpha <= 0 or alpha > 1:
        alpha = 0.1
    diff = utility - running_ema
    return (1 - alpha) * old_var + alpha * (diff * diff)
