# =============================================================================
#  Project   : localvm-research
#  File      : src/localvm/quality/decision_stats.py
#  Purpose   : Shared decision-stability measurement utilities (greedy
#              trajectories, teacher-forced margins/agreement, AUROC, curves)
#  Author    : Simon-Pierre Boucher
#  Contact   : contact@spboucher.ai
#  Created   : 2026-08-12
#  Modified  : 2026-08-12
#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal
#  License   : All rights reserved (research code)
# =============================================================================
"""Decision-stability measurement utilities shared by expG/expD and successors.

First used (inlined) by experiments/micro/expG_decision_stability/benchmark.py
(commit 42c7b3d); extracted here unchanged so later experiments reuse one
implementation. expG's committed copy is kept as-is for reproducibility.
"""

from __future__ import annotations

import mlx.core as mx
import numpy as np


def greedy_generate(model, tokenizer, prompt_ids: list[int], n_tokens: int) -> list[int]:
    """Deterministic greedy generation; returns generated token ids."""
    from mlx_lm.models.cache import make_prompt_cache

    cache = make_prompt_cache(model)
    generated: list[int] = []
    inp = mx.array(list(prompt_ids))[None]
    for _ in range(n_tokens):
        logits = model(inp, cache=cache)
        nxt = int(mx.argmax(logits[0, -1]).item())
        if nxt == tokenizer.eos_token_id:
            break
        generated.append(nxt)
        inp = mx.array([[nxt]])
    return generated


def teacher_forced_stats(model, full_ids: list[int], start: int) -> dict:
    """Forward full_ids once; per-position stats for predictions of tokens
    [start, len). Returns margin (top-1 minus top-2 logit gap), argmax, and
    float16 logprobs (float16 keeps 48 x (128, ~152k) around 2 GB)."""
    logits = model(mx.array(full_ids)[None])[0]
    sel = logits[start - 1 : len(full_ids) - 1].astype(mx.float32)
    top2 = mx.topk(sel, 2, axis=-1)
    argmax = mx.argmax(sel, axis=-1)
    logprobs = sel - mx.logsumexp(sel, axis=-1, keepdims=True)
    mx.eval(top2, argmax, logprobs)
    v = np.array(top2)
    return {
        "margin": np.abs(v[:, 1] - v[:, 0]),
        "argmax": np.array(argmax),
        "logprobs": np.array(logprobs).astype(np.float16),
    }


def auroc(scores: np.ndarray, labels: np.ndarray) -> float:
    """Rank-based AUROC with tie handling (scores: higher = predicted 1)."""
    pos, neg = scores[labels == 1], scores[labels == 0]
    if len(pos) == 0 or len(neg) == 0:
        return float("nan")
    allv = np.concatenate([pos, neg])
    order = np.argsort(allv, kind="mergesort")
    ranks = np.empty(len(order))
    ranks[order] = np.arange(1, len(order) + 1)
    sorted_v = allv[order]
    i = 0
    while i < len(sorted_v):
        j = i
        while j + 1 < len(sorted_v) and sorted_v[j + 1] == sorted_v[i]:
            j += 1
        if j > i:
            ranks[order[i : j + 1]] = ranks[order[i : j + 1]].mean()
        i = j + 1
    r_pos = ranks[: len(pos)].sum()
    return float((r_pos - len(pos) * (len(pos) + 1) / 2) / (len(pos) * len(neg)))


def escalation_curve(margins: np.ndarray, agree: np.ndarray, points: int = 200) -> list[dict]:
    """Escalate tokens with margin < tau (escalated decision assumed exact);
    report escalated fraction vs residual disagreement."""
    qs = np.quantile(margins, np.linspace(0, 1, points))
    out, n = [], len(margins)
    for tau in qs:
        esc = margins < tau
        out.append({
            "tau": float(tau),
            "escalated_frac": float(esc.mean()),
            "residual_disagree": float(np.sum((~esc) & (agree == 0)) / n),
        })
    return out


def kl_ref_vs(model_logprobs_f16: np.ndarray, ref_logprobs_f16: np.ndarray) -> np.ndarray:
    """Per-position KL(ref || model), computed in float32."""
    ref = ref_logprobs_f16.astype(np.float32)
    q = model_logprobs_f16.astype(np.float32)
    return np.sum(np.exp(ref) * (ref - q), axis=-1)