#!/usr/bin/env python3
"""Aggregate MovieBench pixel metrics and prepare/read VBench-5 inputs."""

from __future__ import annotations

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


DIMENSIONS = [
    "subject_consistency", "background_consistency", "motion_smoothness",
    "aesthetic_quality", "imaging_quality",
]
METHODS = ("ffff", "fppf_step2000")


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


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--output_dir", type=Path, required=True)
    args = parser.parse_args()
    rows = []
    for shard in range(4):
        directory = args.output_dir / f"shard_gpu{shard}" / "per_prompt"
        rows.extend(json.loads(path.read_text()) for path in sorted(directory.glob("*.json")))
    rows.sort(key=lambda x: x["prompt_id"])
    ids = [row["prompt_id"] for row in rows]
    if ids != list(range(100)):
        raise RuntimeError(f"Expected MovieBench IDs 0..99, got {ids}")

    mse = [x for row in rows for x in row["mse_per_frame"]]
    ssim = [x for row in rows for x in row["ssim_per_frame"]]
    lpips = [x for row in rows for x in row["lpips_per_frame"]]
    rollout_mse = [x for row in rows for x in row["mse_per_frame"][9:]]
    rollout_ssim = [x for row in rows for x in row["ssim_per_frame"][9:]]
    rollout_lpips = [x for row in rows for x in row["lpips_per_frame"][9:]]
    pixel_metrics = {
        "status": "complete", "checkpoint_step": 2000, "source_layer": 17,
        "num_prompts": 100, "num_frames": len(mse), "prompt_ids": ids,
        "generation_prompts": "MovieGenVideoBench_extended.txt first 100",
        "evaluation_prompts": "MovieGenVideoBench.txt first 100",
        "schedule": "chunk0=FFFF; chunks1-6=FPPF",
        "psnr": -10 * math.log10(sum(mse) / len(mse)),
        "ssim": sum(ssim) / len(ssim), "lpips": sum(lpips) / len(lpips),
        "rollout_psnr": -10 * math.log10(sum(rollout_mse) / len(rollout_mse)),
        "rollout_ssim": sum(rollout_ssim) / len(rollout_ssim),
        "rollout_lpips": sum(rollout_lpips) / len(rollout_lpips),
        "mean_ffff_generation_time_s": sum(x["ffff"]["generation_time_s"] for x in rows) / 100,
        "mean_fppf_generation_time_s": sum(x["fppf"]["generation_time_s"] for x in rows) / 100,
    }
    atomic_json(args.output_dir / "pixel_metrics.json", pixel_metrics)

    for method in METHODS:
        video_dir = args.output_dir / "vbench_inputs" / method
        video_dir.mkdir(parents=True, exist_ok=True)
        info = []
        for row in rows:
            prompt_id = row["prompt_id"]
            matches = list(args.output_dir.glob(f"shard_gpu*/videos/{method}/{prompt_id:05d}.mp4"))
            if len(matches) != 1:
                raise RuntimeError(f"Expected one {method} video for {prompt_id}: {matches}")
            target = video_dir / f"{prompt_id:05d}.mp4"
            if target.exists() or target.is_symlink():
                target.unlink()
            target.symlink_to(matches[0].resolve())
            info.append({
                "video_list": [target.name], "prompt_en": row["original_prompt"],
                "dimension": DIMENSIONS,
            })
        atomic_json(video_dir / "full_info.json", info)

    vbench_rows = []
    for method in METHODS:
        result_dir = args.output_dir / "vbench_results" / method
        matches = list(result_dir.glob("*_eval_results.json"))
        if not matches:
            continue
        if len(matches) != 1:
            raise RuntimeError(f"Expected one VBench result for {method}: {matches}")
        result = json.loads(matches[0].read_text())
        scores = {dimension: float(result[dimension][0]) for dimension in DIMENSIONS}
        vbench_rows.append({"method": method, **scores, "vbench5_mean": sum(scores.values()) / 5})
    if vbench_rows:
        atomic_json(args.output_dir / "vbench_summary.json", vbench_rows)
        with (args.output_dir / "vbench_summary.csv").open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=list(vbench_rows[0]))
            writer.writeheader(); writer.writerows(vbench_rows)
    print(json.dumps({"pixel_metrics": pixel_metrics, "vbench": vbench_rows}, indent=2))


if __name__ == "__main__":
    main()