#!/usr/bin/env python3
"""Freeze evaluation splits so every run is scored on identical states."""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from jevon.data import SamplerConfig, freeze_split


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--out-dir", default="data/frozen")
    p.add_argument("--val-states", type=int, default=180)
    p.add_argument("--test-states", type=int, default=240)
    p.add_argument("--ood-states", type=int, default=160)
    p.add_argument("--seed", type=int, default=20260919)
    args = p.parse_args()

    out = Path(args.out_dir)
    out.mkdir(parents=True, exist_ok=True)
    base = SamplerConfig(augment=False)
    # Sizes never sampled in training; this is the size-generalisation probe.
    ood = SamplerConfig(maze_sizes=(41, 51), snake_sizes=(20, 24), augment=False)

    manifest = {}
    manifest["val"] = freeze_split(out / "val.jsonl", base, "val", args.val_states,
                                   args.seed)
    manifest["test"] = freeze_split(out / "test.jsonl", base, "test",
                                    args.test_states, args.seed + 1)
    manifest["ood"] = freeze_split(out / "ood.jsonl", ood, "test", args.ood_states,
                                   args.seed + 2)
    manifest["config"] = {"train_maze_sizes": list(base.maze_sizes),
                          "train_snake_sizes": list(base.snake_sizes),
                          "ood_maze_sizes": list(ood.maze_sizes),
                          "ood_snake_sizes": list(ood.snake_sizes)}
    (out / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()