"""
Generate the proof-gated planned comparison dataset.

The dataset contains one calibration push per episode and a held-out query
displacement target under a fixed query force.
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np

from planned_comparison_common import (
    OOD_CAMERAS,
    TRAIN_CAMERAS,
    generate_episode,
)


def write_episode(path: Path, episode: dict) -> None:
    np.savez_compressed(
        path,
        rgb_pre=episode["rgb_pre"],
        rgb_post=episode["rgb_post"],
        camera_intrinsics=episode["camera_intrinsics"],
        camera_extrinsics=episode["camera_extrinsics"],
        calibration_force=episode["calibration_force"],
        query_force=episode["query_force"],
        structured_pre=episode["structured_pre"],
        structured_post=episode["structured_post"],
        structured_delta_x=episode["structured_delta_x"],
        structured_mobility=episode["structured_mobility"],
        target_query_dx=episode["target_query_dx"],
        detected_pre_uv=episode["detected_pre_uv"],
        detected_post_uv=episode["detected_post_uv"],
        mass=episode["mass"],
        damping=episode["damping"],
        gt_pre_world=episode["gt_pre_world"],
        gt_post_world=episode["gt_post_world"],
        gt_delta_x=episode["gt_delta_x"],
    )


def generate_split(
    rng: np.random.Generator,
    output_dir: Path,
    split: str,
    camera_name: str,
    count: int,
    starting_index: int,
) -> tuple[list[dict], int]:
    manifest_entries: list[dict] = []
    index = starting_index
    extraction_errors = []

    for _ in range(count):
        episode = generate_episode(rng, camera_name)
        file_name = f"{split}_{camera_name}_{index:05d}.npz"
        write_episode(output_dir / file_name, episode)
        extraction_errors.append(abs(float(episode["structured_delta_x"]) - float(episode["gt_delta_x"])))
        manifest_entries.append(
            {
                "file": file_name,
                "split": split,
                "camera_name": camera_name,
                "material_name": episode["material_name"],
                "calibration_force": float(episode["calibration_force"]),
                "query_force": float(episode["query_force"]),
                "target_query_dx": float(episode["target_query_dx"]),
            }
        )
        index += 1

    if extraction_errors:
        avg_error = float(np.mean(extraction_errors))
        max_error = float(np.max(extraction_errors))
        print(
            f"[generate_planned_comparison_data] split={split} camera={camera_name} "
            f"count={count} avg_state_dx_error={avg_error:.6f} max_state_dx_error={max_error:.6f}"
        )
    return manifest_entries, index


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate planned comparison episodes.")
    parser.add_argument("--output_dir", default="training/data/planned_comparison")
    parser.add_argument("--train_count", type=int, default=300)
    parser.add_argument("--val_count", type=int, default=80)
    parser.add_argument("--test_id_count", type=int, default=80)
    parser.add_argument("--test_ood_per_camera", type=int, default=120)
    parser.add_argument("--seed", type=int, default=0)
    args = parser.parse_args()

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    rng = np.random.default_rng(args.seed)

    manifest: list[dict] = []
    index = 0

    for camera_name in TRAIN_CAMERAS:
        entries, index = generate_split(rng, output_dir, "train", camera_name, args.train_count, index)
        manifest.extend(entries)
        entries, index = generate_split(rng, output_dir, "val", camera_name, args.val_count, index)
        manifest.extend(entries)
        entries, index = generate_split(rng, output_dir, "test_id", camera_name, args.test_id_count, index)
        manifest.extend(entries)

    for camera_name in OOD_CAMERAS:
        entries, index = generate_split(
            rng,
            output_dir,
            "test_ood",
            camera_name,
            args.test_ood_per_camera,
            index,
        )
        manifest.extend(entries)

    with open(output_dir / "manifest.json", "w", encoding="utf-8") as handle:
        json.dump(manifest, handle, indent=2)

    metadata = {
        "seed": args.seed,
        "train_cameras": list(TRAIN_CAMERAS),
        "ood_cameras": list(OOD_CAMERAS),
        "counts": {
            "train_per_camera": args.train_count,
            "val_per_camera": args.val_count,
            "test_id_per_camera": args.test_id_count,
            "test_ood_per_camera": args.test_ood_per_camera,
        },
    }
    with open(output_dir / "metadata.json", "w", encoding="utf-8") as handle:
        json.dump(metadata, handle, indent=2)

    print(f"[generate_planned_comparison_data] wrote {len(manifest)} episodes to {output_dir}")


if __name__ == "__main__":
    main()