"""
Evaluates the model on the harness's own canonical benchmark scrambles.

Separate from the in-training holdout on purpose. The holdout comes from the same
sampler as the training data, so it measures generalization to unseen states but
not robustness to a different generator. These states come from
`generateScrambleSet` -- a different code path, the one every LLM on the board was
actually scored against, at the standard seed = depth * 1000 -- so a gap between
the two numbers would say the model fitted its own sampler's distribution rather
than the task.

    python eval_canonical.py --checkpoint checkpoints/cube
    python eval_canonical.py --distribution-only      # no model needed
"""
import argparse, json, sys
from collections import defaultdict
from pathlib import Path

FIXTURE = Path(__file__).parent / "fixtures" / "canonical_states.json"


def distribution_report(states):
    """Kociemba solution length per depth bucket.

    This is the God's Number plateau made visible: if scramble depth stops being a
    difficulty axis past ~20 moves, then depth 25, 50 and 100 should all show the
    same distance-to-solved as depth 20, and only the shallow buckets should look
    different. It also shows whether the shallow buckets are genuinely easier, or
    whether the scramble generator's redundant moves make a nominal depth-6 state
    closer to solved than 6.
    """
    import kociemba

    by_depth = defaultdict(list)
    for s in states:
        by_depth[s["depth"]].append(len(kociemba.solve(s["facelets"]).split()))

    print(f"{'depth':>6} {'n':>4} {'mean':>7} {'min':>4} {'max':>4}")
    for depth in sorted(by_depth):
        lens = by_depth[depth]
        print(f"{depth:>6} {len(lens):>4} {sum(lens)/len(lens):>7.2f} {min(lens):>4} {max(lens):>4}")


def main():
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--checkpoint", default="")
    p.add_argument("--distribution-only", action="store_true")
    p.add_argument("--limit-per-depth", type=int, default=0)
    p.add_argument("--chunk", type=int, default=0,
                   help="Solve interactively: apply this many moves, then re-show the "
                        "model the true state and ask again. 0 = one-shot (default). "
                        "Interactive results are a different mode -- never table them "
                        "alongside one-shot scores.")
    p.add_argument("--max-rounds", type=int, default=12)
    args = p.parse_args()

    states = json.loads(FIXTURE.read_text())
    if args.limit_per_depth:
        seen = defaultdict(int)
        keep = []
        for s in states:
            if seen[s["depth"]] < args.limit_per_depth:
                keep.append(s)
                seen[s["depth"]] += 1
        states = keep

    if args.distribution_only:
        distribution_report(states)
        return

    if not args.checkpoint:
        p.error("--checkpoint is required unless --distribution-only")

    import torch
    from transformers import LlamaForCausalLM
    from gen_data import SOLVED, apply_sequence
    from serve import solve, solve_chunked

    device = "cuda" if torch.cuda.is_available() else "cpu"
    model = LlamaForCausalLM.from_pretrained(args.checkpoint).to(device).eval()

    mode = f"chunked (chunk={args.chunk}, max_rounds={args.max_rounds})" if args.chunk else "one-shot"
    print(f"mode: {mode}")
    by_depth = defaultdict(lambda: [0, 0, 0])
    for s in states:
        if args.chunk:
            moves = solve_chunked(model, s["facelets"], device,
                                  chunk=args.chunk, max_rounds=args.max_rounds)
        else:
            moves = solve(model, s["facelets"], device).split()
        ok = bool(moves) and apply_sequence(s["facelets"], moves) == SOLVED
        by_depth[s["depth"]][0] += ok
        by_depth[s["depth"]][1] += 1
        by_depth[s["depth"]][2] += len(moves)

    print(f"{'depth':>6} {'solved':>8} {'n':>4} {'rate':>7} {'avg moves':>10}")
    total_ok = total_n = 0
    for depth in sorted(by_depth):
        ok, n, mv = by_depth[depth]
        total_ok, total_n = total_ok + ok, total_n + n
        print(f"{depth:>6} {ok:>8} {n:>4} {ok/n:>6.1%} {mv/n:>10.1f}")
    print(f"{'ALL':>6} {total_ok:>8} {total_n:>4} {total_ok/total_n:>6.1%}")


if __name__ == "__main__":
    main()