#!/usr/bin/env python3
"""Compare two ming_bench.py output dirs (reference vs candidate), stem by stem.

Conditioning (what the DiT receives): cosine similarity over the whole tensor, the
per-token cosine (mean and worst token), and relative L2 = |a - b| / |a|.
Images: MAE, PSNR, windowed 7x7 SSIM on luminance, and alpha MAE for RGBA.

  usage: fidelity_compare.py <reference_dir> <candidate_dir> [--json out.json]
"""
import json
import sys
from pathlib import Path

import numpy as np
from PIL import Image
from safetensors.numpy import load_file


def load_image(path):
    im = Image.open(path)
    rgb = np.asarray(im.convert("RGB"), dtype=np.float64)
    alpha = np.asarray(im.convert("RGBA"), dtype=np.float64)[..., 3] if im.mode in ("RGBA", "LA") else None
    return rgb, alpha, im.size


def box(x, k):
    c = np.cumsum(np.cumsum(np.pad(x, ((1, 0), (1, 0))), 0), 1)
    return (c[k:, k:] - c[:-k, k:] - c[k:, :-k] + c[:-k, :-k]) / (k * k)


def ssim(a, b, k=7, L=255.0):
    c1, c2 = (0.01 * L) ** 2, (0.03 * L) ** 2
    mu_a, mu_b = box(a, k), box(b, k)
    va, vb = box(a * a, k) - mu_a ** 2, box(b * b, k) - mu_b ** 2
    cov = box(a * b, k) - mu_a * mu_b
    s = ((2 * mu_a * mu_b + c1) * (2 * cov + c2)) / ((mu_a ** 2 + mu_b ** 2 + c1) * (va + vb + c2))
    return float(s.mean())


def image_metrics(ref_path, cand_path):
    ra, aa, sa = load_image(ref_path)
    rb, ab, sb = load_image(cand_path)
    if sa != sb:
        raise SystemExit(f"size mismatch {ref_path} {sa} vs {cand_path} {sb}")
    lum = lambda x: 0.299 * x[..., 0] + 0.587 * x[..., 1] + 0.114 * x[..., 2]
    mse = float(((ra - rb) ** 2).mean())
    out = {
        "mae": round(float(np.abs(ra - rb).mean()), 3),
        "psnr_db": None if mse == 0 else round(10 * np.log10(255.0 ** 2 / mse), 2),
        "ssim_lum": round(ssim(lum(ra), lum(rb)), 4),
    }
    if aa is not None and ab is not None:
        out["alpha_mae"] = round(float(np.abs(aa - ab).mean()), 3)
    return out


def cond_metrics(ref_path, cand_path):
    ref, cand = load_file(str(ref_path)), load_file(str(cand_path))
    out = {}
    for key in sorted(set(ref) & set(cand)):
        a, b = ref[key].astype(np.float64), cand[key].astype(np.float64)
        if a.shape != b.shape:
            raise SystemExit(f"{key}: shape mismatch {a.shape} vs {b.shape}")
        fa, fb = a.ravel(), b.ravel()
        tok_a, tok_b = a.reshape(-1, a.shape[-1]), b.reshape(-1, b.shape[-1])
        tok_cos = (tok_a * tok_b).sum(-1) / (np.linalg.norm(tok_a, axis=-1) * np.linalg.norm(tok_b, axis=-1))
        out[key] = {
            "shape": list(a.shape),
            "cosine": round(float(fa @ fb / (np.linalg.norm(fa) * np.linalg.norm(fb))), 6),
            "token_cos_mean": round(float(tok_cos.mean()), 6),
            "token_cos_min": round(float(tok_cos.min()), 6),
            "rel_l2": round(float(np.linalg.norm(fa - fb) / np.linalg.norm(fa)), 6),
        }
    missing = sorted(set(ref) ^ set(cand))
    if missing:
        raise SystemExit(f"conditioning keys present on one side only: {missing}")
    return out


def main():
    ref_dir, cand_dir = Path(sys.argv[1]), Path(sys.argv[2])
    stems = sorted(p.stem for p in ref_dir.glob("*.png") if (cand_dir / p.name).exists())
    if not stems:
        raise SystemExit(f"no common images between {ref_dir} and {cand_dir}")
    rows = []
    for stem in stems:
        row = {"stem": stem, "image": image_metrics(ref_dir / f"{stem}.png", cand_dir / f"{stem}.png")}
        rc, cc = ref_dir / f"{stem}.cond.safetensors", cand_dir / f"{stem}.cond.safetensors"
        if rc.exists() and cc.exists():
            row["cond"] = cond_metrics(rc, cc)
        rows.append(row)
        im = row["image"]
        line = f"{stem:32s} SSIM {im['ssim_lum']:.4f}  PSNR {im['psnr_db']}  MAE {im['mae']:.2f}"
        if "alpha_mae" in im:
            line += f"  aMAE {im['alpha_mae']:.2f}"
        for key, c in row.get("cond", {}).items():
            line += f" | {key[:3]} cos {c['cosine']:.6f} tokmin {c['token_cos_min']:.4f} relL2 {c['rel_l2']:.4f}"
        print(line)
    if "--json" in sys.argv:
        Path(sys.argv[sys.argv.index("--json") + 1]).write_text(json.dumps(rows, indent=2))


if __name__ == "__main__":
    main()