"""Per-draw fork cap (cap_forks): a record with ~1k field rows and generate forks that
repeat its whole known-value block costs at most one trunk pass plus about the room in
fork tokens per step, a fresh random subset each draw, whole-array forks always kept."""

from __future__ import annotations

import random
from collections import Counter
from pathlib import Path
from types import SimpleNamespace

from deg.train import (
    Example,
    Item,
    Source,
    batches,
    cap_decide_forks,
    cap_forks,
    field_of,
    fork_room,
)


def _ex(ctx: int, extract: int = 0, anchors: int = 0, counts: int = 0, gen: int = 0,
        gen_len: int = 16_000, uid: str = "r") -> Example:
    c = SimpleNamespace(ids=tuple(range(ctx)), offsets=tuple((i, i + 1) for i in range(ctx)))
    items = ([Item("extract", tuple(range(80)), ptr=f"/e{i}") for i in range(extract)]
             + [Item("anchor", tuple(range(120)), ptr=f"/a{i}") for i in range(anchors)]
             + [Item("count", tuple(range(110)), ptr=f"/a{i}", aux=True)
                for i in range(counts)]
             + [Item("generate", tuple(range(gen_len)), prompt_len=gen_len - 8, ptr=f"/g{i}")
                for i in range(gen)])
    return Example("deg", uid, c, items)  # type: ignore[arg-type]


def _tok(ex: Example) -> int:
    return sum(len(it.seq) for it in ex.items)


def test_small_record_untouched():
    ex = _ex(500, extract=20)
    assert cap_forks(ex, random.Random(0), 12000) is ex


def test_big_record_bounded_whole_array_forks_kept_pool_untouched():
    # the run-2b record that stalled a step for 20+ min: 7590 context tokens, 1153 rows
    ex = _ex(7590, extract=900, anchors=61, counts=0, gen=176)
    n_items = len(ex.items)
    room = fork_room(ex, 12000)
    assert room == 6000
    whole = sum(len(it.seq) for it in ex.items if it.kind == "anchor")
    biggest = max(len(it.seq) for it in ex.items)
    seen = set()
    for seed in range(50):
        got = cap_forks(ex, random.Random(seed), 12000)
        assert got is not ex and got.ctx is ex.ctx and got.uid == ex.uid
        assert sum(it.kind == "anchor" for it in got.items) == 61
        assert _tok(got) <= room + whole + biggest
        assert got.subsampled == n_items - len(got.items) > 0
        order = [ex.items.index(it) for it in got.items]
        assert order == sorted(order)  # item order kept
        seen.add(tuple(order))
    assert len(ex.items) == n_items and ex.subsampled == 0  # pooled example untouched
    assert len(seen) > 40  # a fresh subset per draw


def test_every_fork_kind_trains_and_short_forks_uniformly():
    ex = _ex(4000, extract=200, anchors=5, counts=5, gen=10)
    hits: Counter = Counter()
    draws = 4000
    for seed in range(draws):
        for it in cap_forks(ex, random.Random(seed), 12000).items:
            hits[it.ptr if it.kind != "count" else "c" + it.ptr] += 1
    assert all(hits[f"/a{i}"] == draws and hits[f"c/a{i}"] == draws for i in range(5))
    gen = [hits[f"/g{i}"] for i in range(10)]
    assert all(g > 0 for g in gen)  # a fork longer than the room still trains
    ext = [hits[f"/e{i}"] for i in range(200)]
    mean = sum(ext) / len(ext)
    assert 0 < mean < draws
    assert max(ext) < 1.25 * mean and min(ext) > 0.75 * mean  # uniform among equals


def test_record_fork_tokens_override():
    ex = _ex(1000, extract=300)
    got = cap_forks(ex, random.Random(0), 12000, record_fork_tokens=800)
    assert _tok(got) <= 800 and len(got.items) == 10


def test_batches_cap_plain_examples_per_draw():
    big = _ex(7590, extract=900, anchors=10, gen=50, uid="big")
    src = Source(Path("/v3/records.jsonl"))
    src.pool = [big]
    gen = batches([src], None, random.Random(0), budget=12000, jev_frac=0.0, max_records=4,
                  draw=2, cap=lambda e, r: cap_forks(e, r, 12000))
    outs = [next(gen) for _ in range(5)]
    for batch in outs:
        assert all(e.uid == "big" and e.subsampled > 0 for e in batch)
        assert all(_tok(e) <= 6000 + 1200 + 16000 for e in batch)
    assert len(big.items) == 960 and big.subsampled == 0


def _decide_ex() -> Example:
    c = SimpleNamespace(ids=tuple(range(300)), offsets=tuple((i, i + 1) for i in range(300)))
    items = ([Item("choice", tuple(range(40)), ptr=f"/rows/{i}/unit") for i in range(50)]
             + [Item("noul", tuple(range(30)), ptr=f"/rows/{i}/audited") for i in range(20)]
             + [Item("choice", tuple(range(40)), ptr="/kind"),
                Item("score", tuple(range(40)), ptr="/risk")]
             + [Item("extract", tuple(range(30)), ptr=f"/rows/{i}/value") for i in range(50)])
    return Example("deg", "d", c, items)  # type: ignore[arg-type]


def test_field_of():
    assert field_of("/rows/12/cells/3/unit") == "/rows/*/cells/*/unit"
    assert field_of("/kind") == "/kind"


def test_decide_forks_capped_per_field_fresh_subset_others_untouched():
    ex = _decide_ex()
    assert cap_decide_forks(ex, random.Random(0), 0) is ex
    seen = set()
    for seed in range(30):
        got = cap_decide_forks(ex, random.Random(seed), 2)
        per = Counter(field_of(it.ptr) for it in got.items if it.kind != "extract")
        assert per == {"/rows/*/unit": 2, "/rows/*/audited": 2, "/kind": 1, "/risk": 1}
        assert sum(it.kind == "extract" for it in got.items) == 50
        assert got.subsampled == 48 + 18
        order = [ex.items.index(it) for it in got.items]
        assert order == sorted(order)
        seen.add(tuple(order))
    assert len(ex.items) == 122 and len(seen) > 20


def test_cap_forks_applies_the_decide_cap_before_the_room():
    ex = _decide_ex()  # fits the room: only the decide cap applies
    got = cap_forks(ex, random.Random(0), 12000, decide_per_field=1)
    assert sum(it.kind in ("choice", "noul", "score") for it in got.items) == 4
    assert cap_forks(ex, random.Random(0), 12000) is ex