#!/usr/bin/env python3
"""Merge independently evaluated two-block FPPP shards into one result set."""

from __future__ import annotations

import argparse
import csv
import json
import os
import shutil
from pathlib import Path
from typing import Any


PROMPT_IDS = list(range(80, 100))
SCHEDULE = "chunk0=FFFF; chunks1-6=FPPP"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--main_dir", type=Path, required=True)
    parser.add_argument("--shard_dir", type=Path, required=True)
    parser.add_argument("--sweep_dir", type=Path, required=True)
    return parser.parse_args()


def atomic_json(path: Path, value: Any) -> None:
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(value, indent=2, ensure_ascii=False) + "\n",
        encoding="utf-8",
    )
    os.replace(temporary, path)


def read_experiment_names(sweep_dir: Path) -> list[str]:
    with (sweep_dir / "summary.csv").open(
        "r", encoding="utf-8", newline=""
    ) as handle:
        return [row["name"] for row in csv.DictReader(handle)]


def merge_pair_directories(main_dir: Path, shard_dir: Path) -> None:
    for source in shard_dir.glob("pair_*"):
        if source.is_dir():
            shutil.copytree(source, main_dir / source.name, dirs_exist_ok=True)


def collect_rows(main_dir: Path, names: list[str]) -> list[dict[str, Any]]:
    rows = []
    for name in names:
        metrics_path = main_dir / name / "metrics.json"
        if not metrics_path.exists():
            raise FileNotFoundError(f"Missing completed metrics: {metrics_path}")
        metrics = json.loads(metrics_path.read_text(encoding="utf-8"))
        if metrics.get("status") != "complete":
            raise RuntimeError(f"Incomplete metrics for {name}")
        if metrics.get("schedule") != SCHEDULE:
            raise RuntimeError(f"Unexpected schedule for {name}: {metrics.get('schedule')}")
        if metrics.get("prompt_ids") != PROMPT_IDS:
            raise RuntimeError(f"Unexpected prompt split for {name}")
        rows.append(
            {
                "name": metrics["name"],
                "source_layer_1": metrics["source_layers"][0],
                "source_layer_2": metrics["source_layers"][1],
                "pair_kind": metrics["pair_kind"],
                "schedule": metrics["schedule"],
                "num_prompts": metrics["num_prompts"],
                "psnr": metrics["psnr"],
                "ssim": metrics["ssim"],
                "lpips": metrics["lpips"],
                "rollout_psnr": metrics["rollout_psnr"],
                "rollout_ssim": metrics["rollout_ssim"],
                "rollout_lpips": metrics["rollout_lpips"],
                "offline_final_val_flow_mse": metrics[
                    "offline_final_val_flow_mse"
                ],
                "mean_generation_time_s": metrics["mean_generation_time_s"],
            }
        )
    rows.sort(key=lambda row: float(row["lpips"]))
    return rows


def write_summary(main_dir: Path, rows: list[dict[str, Any]]) -> None:
    destination = main_dir / "summary.csv"
    temporary = destination.with_suffix(".csv.tmp")
    with temporary.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=list(rows[0]))
        writer.writeheader()
        writer.writerows(rows)
    os.replace(temporary, destination)
    atomic_json(main_dir / "summary.json", rows)


def main() -> None:
    args = parse_args()
    main_dir = args.main_dir.resolve()
    shard_dir = args.shard_dir.resolve()
    sweep_dir = args.sweep_dir.resolve()
    names = read_experiment_names(sweep_dir)
    if len(names) != 43:
        raise RuntimeError(f"Expected 43 experiments, found {len(names)}")
    merge_pair_directories(main_dir, shard_dir)
    rows = collect_rows(main_dir, names)
    write_summary(main_dir, rows)
    atomic_json(
        main_dir / "manifest.json",
        {
            "status": "complete",
            "architecture": "two_block_predictor",
            "prompt_ids": PROMPT_IDS,
            "experiments": names,
            "rollout_schedule": "FPPP",
            "rollout_definition": SCHEDULE,
            "generation_seed_reset_per_prompt": 0,
            "gpu_shards": {"2": 22, "4": 21},
            "merged_shard_dir": str(shard_dir),
        },
    )
    print(f"[merge] {len(rows)} pairs -> {main_dir / 'summary.csv'}", flush=True)


if __name__ == "__main__":
    main()