#!/usr/bin/env python3
"""Merge per-prompt Causal FPPF pixel metrics produced by GPU shards."""

from __future__ import annotations

import argparse
import csv
import json
import math
from pathlib import Path


ROLLOUT_START_FRAME = 9


def mean_std(values: list[float]) -> tuple[float, float]:
    mean = sum(values) / len(values)
    variance = sum((value - mean) ** 2 for value in values) / len(values)
    return mean, math.sqrt(variance)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output_dir", type=Path, required=True)
    parser.add_argument("--experiment", required=True)
    parser.add_argument("--num_shards", type=int, default=4)
    args = parser.parse_args()

    prompt_results = []
    for shard in range(args.num_shards):
        directory = (
            args.output_dir / f"shard_gpu{shard}" / args.experiment / "per_prompt"
        )
        for path in sorted(directory.glob("prompt_*.json")):
            prompt_results.append(json.loads(path.read_text(encoding="utf-8")))
    prompt_results.sort(key=lambda item: int(item["prompt_id"]))
    ids = [int(item["prompt_id"]) for item in prompt_results]
    if ids != list(range(80, 100)):
        raise RuntimeError(f"Expected prompt IDs 80..99, got {ids}")

    mse = [x for item in prompt_results for x in item["mse_per_frame"]]
    psnr = [x for item in prompt_results for x in item["psnr_per_frame"]]
    ssim = [x for item in prompt_results for x in item["ssim_per_frame"]]
    lpips = [x for item in prompt_results for x in item["lpips_per_frame"]]
    rollout_mse = [
        x for item in prompt_results for x in item["mse_per_frame"][ROLLOUT_START_FRAME:]
    ]
    rollout_ssim = [
        x for item in prompt_results for x in item["ssim_per_frame"][ROLLOUT_START_FRAME:]
    ]
    rollout_lpips = [
        x for item in prompt_results for x in item["lpips_per_frame"][ROLLOUT_START_FRAME:]
    ]
    pixel_mse = sum(mse) / len(mse)
    psnr_mean, psnr_std = mean_std(psnr)
    ssim_mean, ssim_std = mean_std(ssim)
    lpips_mean, lpips_std = mean_std(lpips)
    rollout_pixel_mse = sum(rollout_mse) / len(rollout_mse)
    metrics = {
        "status": "complete",
        "name": args.experiment,
        "source_layer": 17,
        "checkpoint_step": 2000,
        "schedule": "chunk0=FFFF; chunks1-6=FPPF",
        "reference": "matching Causal-Forcing FFFF, same prompt and seed",
        "pixel_domain": "VAE-decoded RGB, rounded to uint8",
        "aggregation": "PSNR from global pixel MSE; SSIM/LPIPS mean over decoded frames",
        "num_prompts": len(prompt_results),
        "num_frames": len(psnr),
        "prompt_ids": ids,
        "pixel_mse": pixel_mse,
        "psnr": -10.0 * math.log10(max(pixel_mse, 1e-12)),
        "psnr_frame_mean": psnr_mean,
        "psnr_frame_std": psnr_std,
        "ssim": ssim_mean,
        "ssim_frame_std": ssim_std,
        "lpips": lpips_mean,
        "lpips_frame_std": lpips_std,
        "rollout_start_frame": ROLLOUT_START_FRAME,
        "rollout_num_frames": len(rollout_mse),
        "rollout_pixel_mse": rollout_pixel_mse,
        "rollout_psnr": -10.0 * math.log10(max(rollout_pixel_mse, 1e-12)),
        "rollout_ssim": sum(rollout_ssim) / len(rollout_ssim),
        "rollout_lpips": sum(rollout_lpips) / len(rollout_lpips),
        "mean_generation_time_s": sum(x["generation_time_s"] for x in prompt_results)
        / len(prompt_results),
        "full_calls_per_prompt": prompt_results[0]["full_calls"],
        "predictor_calls_per_prompt": prompt_results[0]["predictor_calls"],
    }
    destination = args.output_dir / args.experiment
    destination.mkdir(parents=True, exist_ok=True)
    (destination / "metrics.json").write_text(
        json.dumps(metrics, indent=2, ensure_ascii=False) + "\n", encoding="utf-8"
    )
    fields = [
        "name", "source_layer", "checkpoint_step", "num_prompts", "num_frames",
        "psnr", "ssim", "lpips", "rollout_psnr", "rollout_ssim", "rollout_lpips",
        "mean_generation_time_s",
    ]
    with (args.output_dir / "summary.csv").open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fields)
        writer.writeheader()
        writer.writerow({key: metrics[key] for key in fields})
    print(json.dumps(metrics, indent=2))


if __name__ == "__main__":
    main()