"""Find real reasoning loops in finished runs, down to the token where the loop starts."""
from __future__ import annotations

import glob
import gzip
import json
import os

import numpy as np

from .protocol import render


def find_loop(ids, max_period: int = 1500, min_span: int = 64, min_reps: int = 3):
    """If the token list ends in one block of `period` tokens repeated back to back (at least `min_reps` times
    and over at least `min_span` tokens), return (onset, period), where onset is the first token of the
    repeating stretch. Otherwise None. The last token is ignored: generation can stop in the middle of a block."""
    a = np.asarray(ids[:-1])
    n = len(a)
    for p in range(1, max_period + 1):
        reps = max(min_reps, -(-min_span // p))
        if reps * p > n:
            break
        if a[n - 1] != a[n - 1 - p]:
            continue
        if all(np.array_equal(a[n - (j + 1) * p:n - j * p], a[n - p:]) for j in range(1, reps)):
            t = n - p - 1
            while t >= 0 and a[t] == a[t + p]:
                t -= 1
            return t + 1, p
    return None


def read_rows(lab: str, data: str) -> dict:
    with open(os.path.join(lab, "data", data + ".jsonl")) as f:
        return {r["id"]: r for r in map(json.loads, f)}


def iter_samples(run_dir: str):
    """Yield (record, j, sample, text) for every sample of a run (finished shards only)."""
    texts = {}
    for f in sorted(glob.glob(os.path.join(run_dir, "shard*.texts.jsonl.gz"))):
        with gzip.open(f, "rt") as fh:
            for d in map(json.loads, fh):
                texts[(d["id"], d["j"])] = d["text"]
    for f in sorted(glob.glob(os.path.join(run_dir, "shard*of*.jsonl"))):
        with open(f) as fh:
            for rec in map(json.loads, fh):
                for j, s in enumerate(rec["samples"]):
                    if (rec["id"], j) in texts:
                        yield rec, j, s, texts[(rec["id"], j)]


def looping_samples(tok, run_dir: str, rows: dict, thinking: bool = True, limit: int | None = None):
    """Every sample of a run that ends in a verbatim token loop: the token ids of prompt + answer (the answer is
    re-tokenized from its text, which almost always gives back the ids the model generated) and where the loop
    starts, counted from the start of the prompt."""
    out = []
    for rec, j, s, text in iter_samples(run_dir):
        if not s.get("cut"):
            continue
        prompt = tok(render(tok, rows[rec["id"]], thinking), add_special_tokens=False).input_ids
        ans = tok(text, add_special_tokens=False).input_ids
        hit = find_loop(ans)
        if hit is None:
            continue
        onset, period = hit
        out.append({"id": rec["id"], "j": j, "ids": prompt + ans, "n_prompt": len(prompt),
                    "onset": len(prompt) + onset, "period": period, "text": text})
        if limit and len(out) >= limit:
            break
    return out