#!/usr/bin/env python3
"""Validate and merge FrameSkip cache shards into one portable cache."""

from __future__ import annotations

import argparse
import json
from pathlib import Path


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--required-draws", required=True)
    parser.add_argument("--shard-dir", required=True)
    parser.add_argument("--shard-count", type=int, required=True)
    parser.add_argument("--output", required=True)
    args = parser.parse_args()

    required = json.loads(Path(args.required_draws).read_text(encoding="utf-8"))
    expected = {
        (dataset_name, str(trajectory_id))
        for dataset_name, trajectory_ids in required["datasets"].items()
        for trajectory_id in trajectory_ids
    }
    records: dict[tuple[str, str], dict] = {}
    shard_paths = []
    for shard_index in range(args.shard_count):
        path = Path(args.shard_dir) / f"shard_{shard_index:02d}_of_{args.shard_count:02d}.jsonl"
        if not path.exists():
            raise RuntimeError(f"missing FrameSkip cache shard: {path}")
        shard_paths.append(path.name)
        with path.open(encoding="utf-8") as source:
            for line_number, line in enumerate(source, 1):
                record = json.loads(line)
                key = (record["dataset_name"], str(record["trajectory_id"]))
                if key in records:
                    raise RuntimeError(f"duplicate FrameSkip cache key: {key}")
                records[key] = record
    if set(records) != expected:
        missing = sorted(expected - set(records))[:20]
        extra = sorted(set(records) - expected)[:20]
        raise RuntimeError(f"cache coverage mismatch: missing={missing}, extra={extra}")

    payload = {
        "schema_version": 1,
        "method": "FrameSkip-adapted (Gaussian-TPI fallback)",
        "upstream_commit": "965becf63e9cd5cc339fbf7fde40d1db47efa67a",
        "retention_ratio": 0.2,
        "importance": {
            "calculator": "GripperAwareImportanceCalculator",
            "avi_weight": 0.6,
            "vac_weight": 0.2,
            "tpi_weight": 0.2,
            "vac_enabled": True,
            "max_vac_frames": 16,
            "visual_encoder": "vit_small_patch14_dinov2.lvd142m",
            "tpi": "Gaussian sigma_squared=0.2 fallback",
            "gripper_boost": 2.0,
        },
        "pruning": {
            "first_last": True,
            "gripper_transition_threshold": 0.1,
            "top_action_change_percentile": 90,
            "min_frames_per_trajectory": 5,
            "mapping": "numpy.searchsorted to the next retained index, clipped at end",
        },
        "schedule": {"pruned_minibatches": 5, "full_minibatches": 1},
        "records": [
            records[key]
            for key in sorted(records, key=lambda item: (item[0], item[1]))
        ],
        "source_shards": shard_paths,
    }
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_text(json.dumps(payload, separators=(",", ":")) + "\n", encoding="utf-8")
    print(json.dumps({
        "output": str(output),
        "records": len(payload["records"]),
        "actual_ratio_min": min(r["actual_ratio"] for r in payload["records"]),
        "actual_ratio_max": max(r["actual_ratio"] for r in payload["records"]),
        "actual_ratio_mean": sum(r["actual_ratio"] for r in payload["records"]) / len(payload["records"]),
    }, indent=2))


if __name__ == "__main__":
    main()