"""Repetition-loop detection for teacher outputs (spec §6 filters; online abort, D-13/D-17).

Both rules measure the share of recent tokens covered by a 16-gram that occurred earlier in the output:
- filter (offline): last 600 tokens, loop if coverage >= 0.5 (the T1.4 FP8 check's rule);
- online abort: last 256 tokens, abort at >= 0.8, checked between generation chunks.
"""
from __future__ import annotations

import numpy as np

FILTER_TAIL, FILTER_THRESHOLD = 600, 0.5
ONLINE_TAIL, ONLINE_THRESHOLD = 256, 0.8


def loop_score(ids: list[int], tail: int = FILTER_TAIL, n: int = 16) -> float:
    """Share of the last `tail` tokens covered by a repeated n-gram: ~1.0 means a repetition loop."""
    t = ids[-tail:]
    seen, covered = {}, np.zeros(len(t), dtype=bool)
    for i in range(len(t) - n + 1):
        g = tuple(t[i:i + n])
        if g in seen:
            covered[i:i + n] = True
        seen.setdefault(g, i)
    return float(covered.mean()) if len(t) else 0.0


def is_loop(ids: list[int]) -> bool:
    return loop_score(ids, FILTER_TAIL) >= FILTER_THRESHOLD


def should_abort(ids: list[int]) -> bool:
    return len(ids) >= ONLINE_TAIL and loop_score(ids, ONLINE_TAIL) >= ONLINE_THRESHOLD


CUT_N, CUT_WINDOW, CUT_THRESHOLD, CUT_MIN = 8, 200, 0.3, 256  # D-37: clean openings of cut-off answers


def clean_prefix_len(ids: list[int], n: int = CUT_N, win: int = CUT_WINDOW, thr: float = CUT_THRESHOLD) -> int:
    """Tokens of an answer before its repetitive tail: the first window of `win` tokens in which n-grams repeated
    from earlier in the answer cover more than `thr` starts the tail. The whole answer if no window does."""
    rep = np.zeros(len(ids), dtype=bool)
    seen = set()
    for i in range(len(ids) - n + 1):
        g = tuple(ids[i:i + n])
        if g in seen:
            rep[i:i + n] = True
        seen.add(g)
    if len(ids) < win:
        return len(ids) if rep.mean() <= thr else 0
    cov = np.convolve(rep.astype(float), np.ones(win) / win, mode="valid")  # cov[j] covers [j, j + win)
    bad = np.nonzero(cov > thr)[0]
    return len(ids) if not len(bad) else int(bad[0])