#!/usr/bin/env python3
"""Merge four independent consecutive three-block training sweeps."""

from __future__ import annotations

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


EXPECTED_TRIPLES = [(index, index + 1, index + 2) for index in range(28)]
EXPECTED_PROMPTS = {"train": list(range(80)), "val": list(range(80, 100))}
GPU_SHARDS = {
    "0": EXPECTED_TRIPLES[0:7],
    "2": EXPECTED_TRIPLES[7:14],
    "6": EXPECTED_TRIPLES[14:21],
    "7": EXPECTED_TRIPLES[21:28],
}


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


def experiment_name(triple: tuple[int, int, int]) -> str:
    return "triple_" + "_".join(f"{layer:02d}" for layer in triple)


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 validate_metrics(
    metrics: dict[str, Any], triple: tuple[int, int, int]
) -> None:
    name = experiment_name(triple)
    expected = {
        "status": "complete",
        "name": name,
        "architecture": "three_block_predictor",
        "initialization_method": "teacher_full",
        "source_layers": list(triple),
        "seed": 0,
        "max_steps": 1000,
        "train_prompt_ids": EXPECTED_PROMPTS["train"],
        "val_prompt_ids": EXPECTED_PROMPTS["val"],
    }
    for key, value in expected.items():
        if metrics.get(key) != value:
            raise RuntimeError(
                f"{name}: expected {key}={value!r}, got {metrics.get(key)!r}"
            )


def result_row(metrics: dict[str, Any]) -> dict[str, Any]:
    return {
        "name": metrics["name"],
        "source_layer_1": metrics["source_layers"][0],
        "source_layer_2": metrics["source_layers"][1],
        "source_layer_3": metrics["source_layers"][2],
        "final_val_flow_mse": metrics["final_val_flow_mse"],
        "final_val_hidden_mse": metrics["final_val_hidden_mse"],
        "final_val_total_loss": metrics["final_val_total_loss"],
        "val_flow_mse_auc": metrics["val_flow_mse_auc"],
        "val_hidden_mse_auc": metrics["val_hidden_mse_auc"],
        "val_total_loss_auc": metrics["val_total_loss_auc"],
        "training_time_s": metrics["training_time_s"],
    }


def write_summary(output_dir: Path, rows: list[dict[str, Any]]) -> None:
    rows.sort(key=lambda row: float(row["final_val_flow_mse"]))
    destination = output_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(output_dir / "summary.json", rows)


def main() -> None:
    args = parse_args()
    shard_root = args.shard_root.resolve()
    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    rows = []
    fingerprints = set()
    for gpu, triples in GPU_SHARDS.items():
        shard_dir = shard_root / f"gpu{gpu}"
        shard_manifest = json.loads(
            (shard_dir / "sweep_manifest.json").read_text(encoding="utf-8")
        )
        if shard_manifest.get("status") != "complete":
            raise RuntimeError(f"GPU {gpu} shard is not complete")
        if shard_manifest.get("triples") != [list(value) for value in triples]:
            raise RuntimeError(f"GPU {gpu} shard contains unexpected triples")
        fingerprints.add(shard_manifest["batch_schedule_sha256"])

        for triple in triples:
            name = experiment_name(triple)
            source = shard_dir / name
            metrics = json.loads(
                (source / "metrics.json").read_text(encoding="utf-8")
            )
            validate_metrics(metrics, triple)
            destination = output_dir / name
            if destination.exists():
                shutil.rmtree(destination)
            shutil.copytree(source, destination)
            rows.append(result_row(metrics))

    if len(rows) != 28:
        raise RuntimeError(f"Expected 28 results, found {len(rows)}")
    if len(fingerprints) != 1:
        raise RuntimeError(f"Batch schedules differ across shards: {fingerprints}")
    write_summary(output_dir, rows)
    atomic_json(
        output_dir / "sweep_manifest.json",
        {
            "status": "complete",
            "architecture": "three_block_predictor",
            "initialization_method": "teacher_full",
            "triples": [list(value) for value in EXPECTED_TRIPLES],
            "num_triples": 28,
            "max_steps": 1000,
            "seed": 0,
            "train_prompt_ids": EXPECTED_PROMPTS["train"],
            "val_prompt_ids": EXPECTED_PROMPTS["val"],
            "batch_schedule_sha256": fingerprints.pop(),
            "gpu_shards": {
                gpu: [list(value) for value in triples]
                for gpu, triples in GPU_SHARDS.items()
            },
            "source_shard_root": str(shard_root),
        },
    )
    print(
        f"[merge] {len(rows)} triples -> {output_dir / 'summary.csv'}",
        flush=True,
    )


if __name__ == "__main__":
    main()