"""Per-generation metrics (blueprint 2.3).

All metrics operate on a probability vector ``p`` (a distribution over the K knowledge
items / alleles). Global and per-region variants are provided; E3 needs the per-region
forms. Every function is a pure function of its inputs.
"""

from __future__ import annotations

import numpy as np


def heterozygosity(p: np.ndarray) -> float:
    """Expected heterozygosity / Simpson diversity ``H = 1 - sum_i p_i^2``.

    The lineage-health metric and the quantity the re-mint gate reads. ``H=0`` at
    fixation (one item), ``H=1-1/K`` at the uniform distribution.

    Args:
        p (np.ndarray): Probability vector over the K items.

    Returns:
        float: Heterozygosity in [0, 1).
    """
    p = np.asarray(p, dtype=float)
    return float(1.0 - np.sum(p * p))


def forward_kl(p_star: np.ndarray, p: np.ndarray, eps: float) -> float:
    """Forward KL to truth ``D_KL(p_star || p) = sum_i p*_i log(p*_i / p_i)``.

    The correct primary collapse metric: it *diverges* when ``p`` drops mass that
    ``p_star`` has, i.e. it explicitly punishes forgetting the improbable (blueprint 2.3;
    reverse KL is deliberately not used). ``p`` is floored at ``eps`` to stay finite.

    Args:
        p_star (np.ndarray): True distribution.
        p (np.ndarray): Current distribution.
        eps (float): Floor applied to ``p`` before the log.

    Returns:
        float: Forward KL divergence (nats).
    """
    p_star = np.asarray(p_star, dtype=float)
    p = np.asarray(p, dtype=float)
    p_floored = np.maximum(p, eps)
    mask = p_star > 0.0  # 0 * log 0 contributes nothing
    return float(np.sum(p_star[mask] * np.log(p_star[mask] / p_floored[mask])))


def tail_mass(p: np.ndarray, tail_mask: np.ndarray) -> float:
    """Total probability mass ``p`` places on the designated tail items.

    The direct measure of collapse: the tail is lost first, so ``tail_mass -> 0`` is the
    signature of drift-driven collapse (blueprint 2.3).

    Args:
        p (np.ndarray): Current distribution.
        tail_mask (np.ndarray): Boolean mask of tail items.

    Returns:
        float: Mass on tail items.
    """
    p = np.asarray(p, dtype=float)
    return float(p[np.asarray(tail_mask, dtype=bool)].sum())


def support_size(p: np.ndarray, eps: float) -> int:
    """Number of items with mass above ``eps`` (surviving items).

    Args:
        p (np.ndarray): Current distribution.
        eps (float): Support threshold.

    Returns:
        int: Count of items with ``p_i > eps``.
    """
    p = np.asarray(p, dtype=float)
    return int(np.sum(p > eps))


def tail_band_metrics(p: np.ndarray, p_star: np.ndarray, tail_mask: np.ndarray,
                      n_bands: int = 4, alive_eps: float = 1e-9):
    """Stratify the tail into ``n_bands`` equal-count rarity bands (band 0 = rarest).

    Makes the per-item survival threshold ``m·p*_i ≳ 1`` (blueprint prediction 4) visible
    band-wise: deeper (rarer) bands sit strictly below shallower ones and the gap narrows
    as grounding rises. Both returned arrays are bounded in [0, 1]; single-run values are
    noisy, so average over replicates before plotting.

    Args:
        p (np.ndarray): Current distribution.
        p_star (np.ndarray): True distribution.
        tail_mask (np.ndarray): Boolean tail mask.
        n_bands (int): Number of equal-count rarity bands.
        alive_eps (float): An item is "alive" if ``p_i > alive_eps``.

    Returns:
        tuple[np.ndarray, np.ndarray]: ``frac_alive[b]`` (fraction of band-b items alive)
        and ``truth_mass_alive[b]`` (share of band-b's TRUE mass carried by alive items),
        each length ``n_bands``.
    """
    p = np.asarray(p, dtype=float)
    p_star = np.asarray(p_star, dtype=float)
    idx = np.where(np.asarray(tail_mask, dtype=bool))[0]
    order = idx[np.argsort(p_star[idx])]           # rarest first
    bands = np.array_split(order, n_bands)
    frac_alive = np.empty(n_bands)
    truth_mass_alive = np.empty(n_bands)
    for b, items in enumerate(bands):
        alive = p[items] > alive_eps
        frac_alive[b] = alive.mean()
        ps = p_star[items]
        truth_mass_alive[b] = ps[alive].sum() / ps.sum() if ps.sum() > 0 else np.nan
    return frac_alive, truth_mass_alive


def per_region(func, p: np.ndarray, regions: np.ndarray, *args) -> dict[int, float]:
    """Apply a metric independently to each region's sub-vector.

    Each region's slice of ``p`` is *not* renormalised; the metric sees the raw masses,
    so per-region ``tail_mass`` and ``support_size`` are directly comparable across
    regions (needed for E3). ``heterozygosity`` per region is therefore a within-``p``
    quantity, documented as such.

    Args:
        func: One of the metric callables above (called as ``func(p_region, *args)``).
        p (np.ndarray): Current distribution.
        regions (np.ndarray): Length-K region index per item.
        *args: Extra positional args forwarded to ``func`` (e.g. eps, or a per-region
            tail mask which is sliced automatically if it is a length-K boolean array).

    Returns:
        dict[int, float]: Region index -> metric value.
    """
    p = np.asarray(p, dtype=float)
    regions = np.asarray(regions)
    out: dict[int, float] = {}
    for r in np.unique(regions):
        sel = regions == r
        sliced_args = tuple(
            a[sel] if isinstance(a, np.ndarray) and a.shape == regions.shape else a
            for a in args
        )
        out[int(r)] = func(p[sel], *sliced_args)
    return out