#!/usr/bin/env python3
"""Measure what eight sequential edits do to an image, per engine.

Reads scripts/nuit-<seed>.png + scripts/chain-<engine>-step<N>.png and writes
apps/site/public/benchmarks/sequential-editing/run.json.

Every number the published benchmark page shows comes from this file. Nothing on
that page is typed by hand — if a metric is not computed here, it is not claimed.

    python3 scripts/bench-metrics.py

Metrics, and what each is for:
  stepDelta        mean |Δ| vs the PREVIOUS step — "did this step change anything"
  driftFromSeed    mean |Δ| vs step 0            — accumulated departure
  outsideMaskDelta mean |Δ| vs previous, OUTSIDE that step's mask (masked chain
                   only) — tests the "I only touch the mask" promise
  insideMaskDelta  same, inside the mask — pairs with the above to show the ratio
  tone             mean R/G/B — catches monotonic warming/darkening
  sharpness        variance of the Laplacian, normalised to step 0 — catches mush
  structureIoU     IoU of Sobel edge maps vs step 0 — catches composition drift

Engines whose output size differs from the seed (Kontext returns 1392x752 for a
1408x768 input) are resized to the seed grid before comparison; that resize is
recorded in the JSON so the page can disclose it.
"""

import json
import os
from pathlib import Path

import numpy as np
from PIL import Image

ROOT = Path(__file__).resolve().parent.parent
SCRIPTS = ROOT / "scripts"
OUT = ROOT / "apps/site/public/benchmarks/sequential-editing"
SRC_DATA = ROOT / "apps/site/src/data/benchmarks"

SEED = SCRIPTS / "nuit-d5bd28cc.png"
STEPS = 8

# Engine → (label, mode). Mask names come from the EDITS array in edit-chain-test.mjs.
ENGINES = [
    ("kontext", "Flux Kontext Pro", "prompt"),
    ("seededit", "SeedEdit 3.0", "prompt"),
    ("recraft", "Recraft inpaint v3", "masked"),
]

MASKS = ["roof", "path", "gable_end", "chimney", "sconces", "door", "deck", "base"]


def load(path, size=None):
    im = Image.open(path).convert("RGB")
    resized = False
    if size and im.size != size:
        im = im.resize(size, Image.LANCZOS)
        resized = True
    return np.asarray(im, dtype=np.float64), resized


def mean_abs_diff(a, b, mask=None):
    d = np.abs(a - b).mean(axis=2)
    if mask is not None:
        if not mask.any():
            return None
        d = d[mask]
    return float(d.mean())


def laplacian_var(a):
    g = a.mean(axis=2)
    k = np.array([[0, 1, 0], [1, -4, 1], [0, 1, 0]], dtype=np.float64)
    h, w = g.shape
    out = np.zeros((h - 2, w - 2))
    for dy in range(3):
        for dx in range(3):
            if k[dy, dx]:
                out += k[dy, dx] * g[dy : dy + h - 2, dx : dx + w - 2]
    return float(out.var())


def edge_map(a, pct=90):
    """Binary edge map at a fixed percentile, so IoU compares structure not contrast."""
    g = a.mean(axis=2)
    h, w = g.shape
    gx = g[1:-1, 2:] - g[1:-1, :-2]
    gy = g[2:, 1:-1] - g[:-2, 1:-1]
    mag = np.hypot(gx, gy)
    return mag > np.percentile(mag, pct)


def iou(a, b):
    inter = np.logical_and(a, b).sum()
    union = np.logical_or(a, b).sum()
    return float(inter / union) if union else 0.0


def main():
    seed_img = Image.open(SEED).convert("RGB")
    size = seed_img.size
    seed = np.asarray(seed_img, dtype=np.float64)
    seed_edges = edge_map(seed)
    seed_sharp = laplacian_var(seed)

    # Masks are authored at seed resolution; True = the region the model may change.
    masks = {}
    for m in MASKS:
        p = SCRIPTS / f"mask-{m}.png"
        if p.exists():
            mi = Image.open(p).convert("L")
            if mi.size != size:
                mi = mi.resize(size, Image.NEAREST)
            masks[m] = np.asarray(mi) > 127

    result = {
        "seedImage": str(SEED.relative_to(ROOT)),
        "resolution": f"{size[0]}x{size[1]}",
        "steps": STEPS,
        "harness": "scripts/edit-chain-test.mjs",
        "metricsScript": "scripts/bench-metrics.py",
        "engines": [],
    }

    for key, label, mode in ENGINES:
        frames = [seed]
        resized_any = False
        missing = []
        for n in range(1, STEPS + 1):
            p = SCRIPTS / f"chain-{key}-step{n}.png"
            if not p.exists():
                missing.append(n)
                frames.append(frames[-1])
                continue
            arr, r = load(p, size)
            resized_any = resized_any or r
            frames.append(arr)

        steps = []
        for n in range(1, STEPS + 1):
            cur, prev = frames[n], frames[n - 1]
            mask = masks.get(MASKS[n - 1]) if mode == "masked" else None
            row = {
                "step": n,
                "edit": MASKS[n - 1],
                "stepDelta": round(mean_abs_diff(cur, prev), 3),
                "driftFromSeed": round(mean_abs_diff(cur, seed), 3),
                "tone": [round(float(cur[:, :, c].mean()), 2) for c in range(3)],
                "sharpness": round(laplacian_var(cur) / seed_sharp, 4),
                "structureIoU": round(iou(edge_map(cur), seed_edges), 4),
            }
            if mask is not None:
                outside = mean_abs_diff(cur, prev, ~mask)
                inside = mean_abs_diff(cur, prev, mask)
                row["outsideMaskDelta"] = round(outside, 3) if outside is not None else None
                row["insideMaskDelta"] = round(inside, 3) if inside is not None else None
                # Share of pixels outside the mask that moved by more than 2/255 —
                # the number that disproved our own "pixel-identical" claim.
                d = np.abs(cur - prev).mean(axis=2)[~mask]
                row["outsideMaskChangedPct"] = round(float((d > 2).mean() * 100), 2)
                row["maskCoveragePct"] = round(float(mask.mean() * 100), 2)
            steps.append(row)

        engine = {
            "key": key,
            "label": label,
            "mode": mode,
            "resizedToSeedGrid": resized_any,
            "missingSteps": missing,
            "seedTone": [round(float(seed[:, :, c].mean()), 2) for c in range(3)],
            "steps": steps,
        }

        # Cumulative view: final frame vs the seed, split by the UNION of all eight
        # masks. Per-step and cumulative answer different questions and must never
        # be quoted interchangeably — "did one edit leak" vs "did the whole chain".
        if mode == "masked" and masks:
            union = np.zeros(seed.shape[:2], bool)
            for m in MASKS:
                if m in masks:
                    union |= masks[m]
            final = frames[STEPS]
            d = np.abs(final - seed).mean(axis=2)
            out = d[~union]
            engine["cumulative"] = {
                "unionMaskCoveragePct": round(float(union.mean() * 100), 2),
                "outsideMeanDelta": round(float(out.mean()), 2),
                "insideMeanDelta": round(float(d[union].mean()), 2),
                "outsideChangedPct": {
                    f">{t}": round(float((out > t).mean() * 100), 1) for t in (1, 2, 4, 8)
                },
            }

        result["engines"].append(engine)

    payload = json.dumps(result, indent=2) + "\n"

    # Two destinations, on purpose:
    #   public/  — the citable download the Dataset JSON-LD points at
    #   src/data — the build-time import the page renders from, so no number on
    #              the page can drift from the measured file
    OUT.mkdir(parents=True, exist_ok=True)
    (OUT / "run.json").write_text(payload)
    SRC_DATA.mkdir(parents=True, exist_ok=True)
    (SRC_DATA / "sequential-editing.json").write_text(payload)
    print(f"wrote {OUT / 'run.json'}")
    print(f"wrote {SRC_DATA / 'sequential-editing.json'}")

    for e in result["engines"]:
        s = e["steps"]
        print(f"\n{e['label']} ({e['mode']}){'  [resized]' if e['resizedToSeedGrid'] else ''}")
        print(f"  drift  {s[0]['driftFromSeed']:6.2f} -> {s[-1]['driftFromSeed']:6.2f}")
        print(f"  sharp  {s[0]['sharpness']:6.3f} -> {s[-1]['sharpness']:6.3f}")
        print(f"  IoU    {s[0]['structureIoU']:6.3f} -> {s[-1]['structureIoU']:6.3f}")
        print(f"  tone R {e['seedTone'][0]:6.2f} -> {s[-1]['tone'][0]:6.2f}")
        if "outsideMaskChangedPct" in s[-1]:
            avg = sum(x["outsideMaskChangedPct"] for x in s) / len(s)
            print(f"  outside-mask pixels changed >2: avg {avg:.1f}%")


if __name__ == "__main__":
    main()
