"""
Generates training data from the states the model's *own* policy visits.

## Why

The model is trained only on states drawn from `gen_data.py`'s sampler, each
paired with one Kociemba solution. It is never trained on the states its own
attempts produce -- and measurement shows those are exactly where it is worst.
Tracking distance-to-solved across rounds of interactive solving, from a state 10
moves from solved:

    10 -> 16 -> 17 -> 15 -> 17 -> 17 -> 18 -> 17 -> 18

It walks away from solved and stays there. Those visited states are off the
training distribution, so the model has never been shown what to do in them.

This is textbook exposure bias, and the fix is DAgger's: run the current policy,
collect the states it actually reaches, label those with the expert (Kociemba),
and train on them alongside the original data. Behaviour cloning alone compounds
error quadratically in horizon length; adding the policy's own state
distribution brings it back to linear.

Three other explanations have already been tested and eliminated -- more
parameters, more layers at matched parameters, and interactive chunking all
changed nothing -- which is what makes this the remaining candidate rather than
one option among several.

## How

Rollouts are batched on the GPU: a batch of states is decoded together, the first
`--chunk` moves of each proposal are applied, and the resulting states are
recorded. Labelling is the slow part and runs across all cores, exactly as in
gen_data.py.

    python gen_rollout_data.py --checkpoint <hub-id-or-dir> --count 2000000 \\
        --out data/rollout.jsonl
"""
import argparse, json, random, sys, time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))
from gen_data import (SOLVED, apply_sequence, random_state, effective_cpus,
                      check_kociemba_is_native)
import cube_tokenizer as T


def collect_states(model, device, n_states, chunk, rounds, batch_size, rng, deep_depth,
                   shallow_frac, max_shallow):
    """Rolls the policy out and returns the states it lands in.

    The start states are excluded: those are already in the training
    distribution. Only what the policy *reaches* is new information.
    """
    import torch
    from serve import _decode

    visited = []
    while len(visited) < n_states:
        batch = []
        for _ in range(batch_size):
            depth = rng.randint(1, max_shallow) if rng.random() < shallow_frac else deep_depth
            batch.append(random_state(rng, depth))

        for _ in range(rounds):
            live = [s for s in batch if s != SOLVED]
            if not live:
                break
            with torch.no_grad():
                prompts = torch.tensor([T.encode_state(s) for s in live], device=device)
                out = prompts
                finished = torch.zeros(len(live), dtype=torch.bool, device=device)
                for _ in range(T.MAX_SOLUTION + 1):
                    nxt = model(input_ids=out).logits[:, -1, :].argmax(-1)
                    nxt = torch.where(finished, torch.full_like(nxt, T.PAD), nxt)
                    finished |= nxt == T.EOS
                    out = torch.cat([out, nxt[:, None]], dim=1)
                    if bool(finished.all()):
                        break
                proposals = [T.decode_solution(r) for r in out[:, prompts.shape[1]:].tolist()]

            nxt_batch = []
            for state, moves in zip(live, proposals):
                if not moves:
                    continue
                landed = apply_sequence(state, moves[:chunk])
                if landed != SOLVED:
                    visited.append(landed)
                    nxt_batch.append(landed)
            batch = nxt_batch
            if not batch:
                break
        print(f"  collected {len(visited)}/{n_states}", file=sys.stderr, flush=True)
    return visited[:n_states]


def _label(state):
    import kociemba
    try:
        return state, kociemba.solve(state).split()
    except Exception:
        return state, None


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--count", type=int, default=2_000_000)
    p.add_argument("--chunk", type=int, default=8)
    p.add_argument("--rounds", type=int, default=4)
    p.add_argument("--batch-size", type=int, default=512)
    p.add_argument("--workers", type=int, default=0)
    p.add_argument("--seed", type=int, default=7)
    p.add_argument("--deep-depth", type=int, default=200)
    p.add_argument("--shallow-frac", type=float, default=0.15)
    p.add_argument("--max-shallow", type=int, default=8)
    p.add_argument("--out", required=True)
    args = p.parse_args()

    ms = check_kociemba_is_native()
    workers = args.workers or effective_cpus()
    print(f"kociemba {ms:.1f} ms/solve | workers {workers}", file=sys.stderr)

    import torch
    from transformers import LlamaForCausalLM
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = LlamaForCausalLM.from_pretrained(args.checkpoint).to(device).eval()
    print(f"rolling out {args.checkpoint} on {device}", file=sys.stderr)

    start = time.time()
    rng = random.Random(args.seed)
    states = collect_states(model, device, args.count, args.chunk, args.rounds,
                            args.batch_size, rng, args.deep_depth, args.shallow_frac,
                            args.max_shallow)
    print(f"collected {len(states)} states in {time.time()-start:.0f}s", file=sys.stderr)

    import multiprocessing as mp
    start = time.time()
    written = 0
    with mp.Pool(workers) as pool, open(args.out, "w") as fh:
        for state, moves in pool.imap_unordered(_label, states, chunksize=64):
            if not moves:
                continue
            fh.write(json.dumps({"state": state, "solution": " ".join(moves),
                                 "depth": len(moves)}) + "\n")
            written += 1
    secs = time.time() - start
    print(f"labelled and wrote {written} pairs in {secs:.0f}s ({written/max(secs,1e-9):.0f}/s)",
          file=sys.stderr)


if __name__ == "__main__":
    main()