#!/usr/bin/env python3
"""Freeze the exact FrameSkip cache coverage required by the 20K->20.5K run."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from starVLA.dataloader.infoflow_sampler import LogicalDraw, ValidAnchorManifest


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--manifest", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--start-step", type=int, default=20_000)
    parser.add_argument("--updates", type=int, default=500)
    parser.add_argument("--global-batch-size", type=int, default=64)
    parser.add_argument("--global-microbatch-size", type=int, default=8)
    parser.add_argument("--pruned-batches", type=int, default=5)
    parser.add_argument("--full-batches", type=int, default=1)
    args = parser.parse_args()

    manifest = ValidAnchorManifest.load(args.manifest)
    start_draw = args.start_step * args.global_batch_size
    end_draw = (args.start_step + args.updates) * args.global_batch_size
    period = args.pruned_batches + args.full_batches
    required: dict[str, set[int | str]] = {
        dataset.name: set() for dataset in manifest.datasets
    }
    pruned_draws = 0
    full_draws = 0
    for global_draw_id in range(start_draw, end_draw):
        microbatch = (global_draw_id - start_draw) // args.global_microbatch_size
        is_pruned = microbatch % period < args.pruned_batches
        trace = manifest.sample(args.seed, LogicalDraw(0, global_draw_id))
        if is_pruned:
            required[trace.dataset_name].add(trace.trajectory_id)
            pruned_draws += 1
        else:
            full_draws += 1

    payload = {
        "schema_version": 1,
        "method": "FrameSkip-adapted (Gaussian-TPI fallback)",
        "manifest_path": str(Path(args.manifest).resolve()),
        "sampling_seed": args.seed,
        "start_step": args.start_step,
        "updates": args.updates,
        "global_batch_size": args.global_batch_size,
        "global_microbatch_size": args.global_microbatch_size,
        "start_global_draw_id": start_draw,
        "end_global_draw_id_exclusive": end_draw,
        "pruned_batches": args.pruned_batches,
        "full_batches": args.full_batches,
        "pruned_draws": pruned_draws,
        "full_draws": full_draws,
        "datasets": {
            name: sorted(ids, key=lambda value: (type(value).__name__, str(value)))
            for name, ids in required.items()
        },
    }
    payload["required_trajectory_count"] = sum(
        len(values) for values in payload["datasets"].values()
    )
    destination = Path(args.output)
    destination.parent.mkdir(parents=True, exist_ok=True)
    destination.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    print(json.dumps({
        "output": str(destination),
        "pruned_draws": pruned_draws,
        "full_draws": full_draws,
        "required_trajectory_count": payload["required_trajectory_count"],
        "per_dataset": {k: len(v) for k, v in payload["datasets"].items()},
    }, indent=2))


if __name__ == "__main__":
    main()