#!/usr/bin/env python
"""Issue 5 probe: PSNR of the int8 artifact's vLLM outputs against the
sealed Diffusers baseline attempt 2026-10-07T170729Z-q3-4r.

Usage: psnr.py --fixture FIXTURE_ID RUN_DIR [RUN_DIR ...]

Compares each RUN_DIR/matrix/FIXTURE_ID.sSEED.png against the baseline
sample with the same fixture and seed. RGBA outputs are composited over
white, matching issue 1's comparator. The sealed baseline covers the 1-,
2-, and 4-reference fixtures; the derived 3-reference fixture has no
sealed row (gate.py + visual check cover it).
"""

from __future__ import annotations

import argparse
import json
import math
import pathlib

import numpy as np
from PIL import Image

ATTEMPT = (
    pathlib.Path.home()
    / ".local/share/multi-model-helper-server/data/evaluation/media/image/attempts/2026-10-07T170729Z-q3-4r"
)


def to_rgb(array: np.ndarray) -> np.ndarray:
    if array.shape[2] == 4:
        alpha = array[..., 3:4].astype(np.float64) / 255.0
        bg = np.full_like(array[..., :3], 255, dtype=np.float64)
        return np.round(
            array[..., :3].astype(np.float64) * alpha + bg * (1 - alpha)
        ).clip(0, 255).astype(np.uint8)
    return array[..., :3]


def psnr(a: np.ndarray, b: np.ndarray) -> float:
    mse = float(np.mean((a.astype(np.float64) - b.astype(np.float64)) ** 2))
    return float("inf") if mse == 0 else 10 * math.log10(255.0**2 / mse)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--fixture", default="edit-compose-four")
    parser.add_argument("runs", nargs="+", type=pathlib.Path)
    args = parser.parse_args()
    fixture = args.fixture

    samples = json.loads((ATTEMPT / "samples.json").read_text())["samples"]
    baseline = {
        s["seed"]: ATTEMPT / s["file"]
        for s in samples
        if s["fixture_id"] == fixture and s["status"] == "succeeded"
    }
    summary = {}
    for run_dir in args.runs:
        rows = []
        for seed, base_path in sorted(baseline.items()):
            ours = run_dir / "matrix" / f"{fixture}.s{seed}.png"
            if not ours.exists():
                rows.append({"seed": seed, "error": "missing"})
                continue
            bm, om = Image.open(base_path), Image.open(ours)
            ba, oa = np.asarray(bm), np.asarray(om)
            if ba.shape[:2] != oa.shape[:2]:
                rows.append(
                    {"seed": seed, "error": f"size {oa.shape[:2]} vs {ba.shape[:2]}"}
                )
                continue
            use_rgba = bm.mode == "RGBA" and om.mode == "RGBA"
            a, b = (ba, oa) if use_rgba else (to_rgb(ba), to_rgb(oa))
            rows.append(
                {
                    "seed": seed,
                    "psnr_db": round(psnr(a, b), 2),
                    "modes": f"{bm.mode}/{om.mode}",
                }
            )
        vals = [r["psnr_db"] for r in rows if "psnr_db" in r]
        if vals:
            summary[run_dir.name] = {
                "rows": rows,
                "mean_db": round(sum(vals) / len(vals), 2),
                "min_db": min(vals),
                "max_db": max(vals),
                "n": len(vals),
            }
        else:
            summary[run_dir.name] = {"rows": rows}
    print(json.dumps(summary, indent=1))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())