"""
Diversity math over an embedding group. Pure numpy, no torch.

Given L2-normalized embeddings E = [e_1..e_G] (shape G x d) for ONE prompt's
generation group, we expose two families of per-sample diversity credit:

  1. deviation      d_i = mean_{j!=i} (1 - cos(e_i, e_j))
                    -- pairwise, cheap, but blind to cluster structure:
                    two tight far-apart clusters score as high as a spread.

  2. marginal       m_i = logdet(L) - logdet(L_{-i})
     contribution   -- leave-one-out volume credit under the cosine kernel.
                    A duplicate contributes ~nothing (its direction is already
                    spanned), so this DOES see cluster structure.

Both are per-sample. This is deliberate: a set-level scalar (e.g. logdet(L)
itself) is constant within a GRPO group, so its within-group std is 0 and the
normalized advantage is identically 0. It cannot train anything. Every
quantity here varies across i within a group.
"""

from __future__ import annotations

import numpy as np

EPS_JITTER = 1e-3


def l2_normalize(E: np.ndarray, axis: int = -1) -> np.ndarray:
    """Row-normalize; zero rows are left as zeros rather than NaN."""
    E = np.asarray(E, dtype=np.float64)
    n = np.linalg.norm(E, axis=axis, keepdims=True)
    return E / np.maximum(n, 1e-12)


def cosine_kernel(E: np.ndarray, eps: float = EPS_JITTER) -> np.ndarray:
    """L = E E^T + eps*I.

    E must be L2-normalized, so diag(E E^T) = 1 and L is a correlation-like
    PSD matrix. The eps jitter keeps logdet finite when rows are collinear
    (exact duplicates make E E^T singular).
    """
    E = np.asarray(E, dtype=np.float64)
    L = E @ E.T
    L = 0.5 * (L + L.T)  # kill float asymmetry before Cholesky
    return L + eps * np.eye(L.shape[0], dtype=np.float64)


def pairwise_deviation(E: np.ndarray) -> np.ndarray:
    """d_i = mean_{j != i} (1 - cos(e_i, e_j)). Shape (G,).

    G == 1 has no off-diagonal terms; we return 0.0 (a lone sample has no
    measurable deviation, and 0 is the neutral value for the reward).
    """
    E = np.asarray(E, dtype=np.float64)
    G = E.shape[0]
    if G < 2:
        return np.zeros(G, dtype=np.float64)
    S = E @ E.T
    D = 1.0 - S
    np.fill_diagonal(D, 0.0)
    return D.sum(axis=1) / (G - 1)


def _logdet_psd(L: np.ndarray) -> float:
    """logdet via Cholesky; falls back to slogdet if L drifts non-PD."""
    try:
        c = np.linalg.cholesky(L)
        return float(2.0 * np.sum(np.log(np.diag(c))))
    except np.linalg.LinAlgError:
        sign, ld = np.linalg.slogdet(L)
        if sign <= 0:
            return float("-inf")
        return float(ld)


def logdet_volume(E: np.ndarray, eps: float = EPS_JITTER) -> float:
    """D(Y) = logdet(E E^T + eps I). Set-level scalar.

    Report this as a METRIC. Never hand it to GRPO as a per-sample reward:
    it is identical for every i in the group -> zero advantage.
    """
    if np.asarray(E).shape[0] == 0:
        return 0.0
    return _logdet_psd(cosine_kernel(E, eps))


def marginal_contributions(E: np.ndarray, eps: float = EPS_JITTER) -> np.ndarray:
    """m_i = logdet(L) - logdet(L_{-i}), shape (G,).

    Computed by G explicit leave-one-out logdets. G <= 16 here, so this is
    ~16 Cholesky calls on a 15x15 matrix -- utterly negligible next to a
    single LLM forward pass. Rank-one downdate machinery would be a
    micro-optimization with real numerical-stability downside; not worth it.

    Interpretation: m_i is the log-volume the group loses by dropping i.
    A duplicate has m_i ~ log(eps) -> large negative. An orthogonal direction
    has m_i ~ log(1+eps) ~ 0. So m_i is a *penalty scale*: higher (closer to
    0) means "this sample carries a direction nothing else covers".
    """
    E = np.asarray(E, dtype=np.float64)
    G = E.shape[0]
    if G < 2:
        return np.zeros(G, dtype=np.float64)
    L = cosine_kernel(E, eps)
    full = _logdet_psd(L)
    out = np.empty(G, dtype=np.float64)
    idx = np.arange(G)
    for i in range(G):
        keep = idx[idx != i]
        out[i] = full - _logdet_psd(L[np.ix_(keep, keep)])
    return out


def effective_rank(E: np.ndarray, eps: float = 1e-12) -> float:
    """exp(Shannon entropy of the normalized Gram spectrum). Roy & Vetterli.

    Reads as "how many distinct directions does this set effectively span":
    1.0 for identical stories, G for mutually orthogonal ones, ~2 for two tight
    clusters. Reported as a METRIC, never used as a reward (it is set-level, so
    it is constant within a group and would produce zero advantage).

    This exists because silhouette-selected k-means turned out to be unusable as
    a mode counter on this data: a fully collapsed set and a fully spread set
    both peak near silhouette 0.20 (at k=4 and k=7 respectively), because
    k-means partitions isotropic data happily regardless of spread. Effective
    rank has no k to select and degrades gracefully.
    """
    E = np.asarray(E, dtype=np.float64)
    if E.shape[0] == 0:
        return 0.0
    if E.shape[0] == 1:
        return 1.0
    L = E @ E.T
    w = np.linalg.eigvalsh(0.5 * (L + L.T))
    w = np.clip(w, 0.0, None)
    s = w.sum()
    if s <= 0:
        return 1.0
    p = w / s
    p = p[p > eps]
    return float(np.exp(-(p * np.log(p)).sum()))


def zscore(x: np.ndarray, ddof: int = 0) -> np.ndarray:
    """Per-group z-score. Constant input -> all zeros (not NaN).

    Used on m_i before it enters the reward: raw m_i lives on a log scale
    with a long negative tail (a duplicate pair can hit log(1e-3) ~ -6.9),
    which would otherwise dominate the quality term.
    """
    x = np.asarray(x, dtype=np.float64)
    if x.size == 0:
        return x
    s = x.std(ddof=ddof)
    if not np.isfinite(s) or s < 1e-12:
        return np.zeros_like(x)
    return (x - x.mean()) / s


def scale_to_reference(x: np.ndarray, ref_scale: float = 1.0) -> np.ndarray:
    """Rescale x to have unit-ish spread times ref_scale, preserving mean 0.

    Deviation d_i lives in [0, 2] but in practice clusters in [0.1, 0.5] for
    same-prompt stories -- an order of magnitude below judge quality (0-10).
    Feeding raw d_i with alpha=0.5 would make the diversity term invisible.
    """
    return zscore(x) * float(ref_scale)


def greedy_diverse_subset(
    quality: np.ndarray,
    E: np.ndarray,
    k: int,
    lam: float = 1.0,
    eps: float = EPS_JITTER,
) -> list[int]:
    """Greedily pick k indices maximizing sum(quality) + lam * logdet(L_S).

    Standard submodular greedy: the objective is monotone-ish and DPP logdet
    is submodular, so greedy carries the usual (1 - 1/e) flavor of guarantee.
    Used to build E3's multi-positive chosen sets.

    Returns indices into the rows of E, in selection order.
    """
    quality = np.asarray(quality, dtype=np.float64)
    E = np.asarray(E, dtype=np.float64)
    G = E.shape[0]
    k = int(min(k, G))
    if k <= 0:
        return []
    L = cosine_kernel(E, eps)

    chosen: list[int] = []
    for _ in range(k):
        best_gain, best_i = -np.inf, -1
        for i in range(G):
            if i in chosen:
                continue
            cand = chosen + [i]
            vol = _logdet_psd(L[np.ix_(cand, cand)])
            gain = quality[cand].sum() + lam * vol
            if gain > best_gain:
                best_gain, best_i = gain, i
        if best_i < 0:
            break
        chosen.append(best_i)
    return chosen


def group_metrics(E: np.ndarray, eps: float = EPS_JITTER) -> dict:
    """Bundle of reporting metrics for one group. Metrics only, not rewards."""
    E = np.asarray(E, dtype=np.float64)
    G = E.shape[0]
    if G < 2:
        return {"n": G, "mean_pairwise_dist": 0.0, "logdet": 0.0,
                "mean_marginal": 0.0, "min_marginal": 0.0}
    d = pairwise_deviation(E)
    m = marginal_contributions(E, eps)
    return {
        "n": G,
        "mean_pairwise_dist": float(d.mean()),
        "logdet": float(logdet_volume(E, eps)),
        "mean_marginal": float(m.mean()),
        "min_marginal": float(m.min()),
    }