#!/usr/bin/env python3
"""Iterate a building against the reference photography and the engine render.

The loop this tool implements:

```
reference photograph ─┐
                      ├─ compare ─► metrics + feedback ─► blender-skill actions
engine screenshot    ─┘                    │
                                           └─ next iteration (uncommitted)
```

It answers three questions with numbers instead of opinions:

1. **Does the shot show what it was supposed to?** (`--verify`) — every capture
   has an expectation in `Tools/visual_targets.json`; a shot that is blank, the
   wrong size, or wildly unlike its expectation is reported as a failure.
2. **How far is the render from the reference?** (`--compare`) — luminance,
   contrast, edge density, palette and 3×3 composition are measured on both and
   reported as deltas with a 0–1 similarity score.
3. **What should be modelled next?** (`--feedback`) — the deltas are translated
   into concrete actions phrased for the Blender skill packs (silhouette, relief,
   value range, hue, framing).

Iterations are **not committed**: every artefact goes under
`Art/reference/iterations/<iteration>/`, which is gitignored. Promote a shot into
`Screenshots/` only after the owner confirms it.

    python3 Tools/visual_loop.py --verify                    # gate the current set
    python3 Tools/visual_loop.py --compare --iteration 3     # measure vs references
    python3 Tools/visual_loop.py --compare --iteration 3 --judge   # + vision advisory

Metrics are computed with Pillow (already in requirements-tools.txt); no numpy.
"""

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys

from PIL import Image, ImageFilter, ImageStat

ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
TARGETS = os.path.join(ROOT, "Tools", "visual_targets.json")
ITER_ROOT = os.path.join(ROOT, "Art", "reference", "iterations")

# A shot must clear these floors or it is treated as broken, not merely ugly.
MIN_EDGE_DENSITY = 0.010
MIN_CONTRAST = 12.0
MIN_MEAN_LUMA = 8.0


# --- metrics ----------------------------------------------------------------


def _luma_gray(image):
    return image.convert("L")


def metrics(image_path):
    """Measure one image. Everything is on the 8-bit luma/RGB 0..255 scale."""
    with Image.open(image_path) as handle:
        image = handle.convert("RGB")
    width, height = image.size
    gray = _luma_gray(image)
    stat = ImageStat.Stat(gray)

    # Edge density: fraction of pixels above a low threshold after FIND_EDGES.
    # Stands in for "how much relief and silhouette detail is in the frame".
    # Filter, then crop: FIND_EDGES samples outside the image as zero, so it leaves
    # its own border ring at a high value (236 of 3600 px on a 60x60 frame, 6.5%),
    # which would make a perfectly flat picture look detailed. Cropping *before*
    # filtering does not help — the new border gets the same treatment.
    filtered = gray.filter(ImageFilter.FIND_EDGES)
    inner = filtered.crop((2, 2, max(3, width - 2), max(3, height - 2)))
    hist = inner.histogram()
    area = float(inner.width * inner.height)
    edge_pixels = sum(hist[24:])
    edge_density = edge_pixels / area if area else 0.0

    # Palette: quantise and keep the dominant colours with their share.
    total = float(width * height)
    quantised = image.quantize(colors=12, method=Image.Quantize.MEDIANCUT)
    palette = quantised.getpalette() or []
    counts = sorted(quantised.getcolors() or [], reverse=True)
    dominant = []
    for count, index in counts[:5]:
        base = index * 3
        if base + 2 < len(palette):
            dominant.append({
                "rgb": palette[base:base + 3],
                "share": round(count / total, 3),
            })

    # 3x3 composition grid of mean luma, so "too much sky" is measurable.
    grid = []
    for row in range(3):
        for col in range(3):
            box = (int(col * width / 3), int(row * height / 3),
                   int((col + 1) * width / 3), int((row + 1) * height / 3))
            grid.append(round(ImageStat.Stat(gray.crop(box)).mean[0], 1))

    return {
        "path": os.path.relpath(image_path, ROOT),
        "size": [width, height],
        "mean_luma": round(stat.mean[0], 1),
        "rms_contrast": round(stat.stddev[0], 1),
        "edge_density": round(edge_density, 4),
        "dominant": dominant,
        "grid_luma": grid,
    }


def _palette_distance(a, b):
    """Mean nearest-colour distance between two dominant palettes, 0..441."""
    if not a or not b:
        return None
    total = 0.0
    for entry in a:
        pa = entry["rgb"]
        best = min(
            sum((pa[i] - pb["rgb"][i]) ** 2 for i in range(3)) ** 0.5 for pb in b)
        total += best
    return round(total / len(a), 1)


def _correlation(values_a, values_b):
    n = len(values_a)
    mean_a = sum(values_a) / n
    mean_b = sum(values_b) / n
    var_a = sum((v - mean_a) ** 2 for v in values_a)
    var_b = sum((v - mean_b) ** 2 for v in values_b)
    if var_a <= 0.0 or var_b <= 0.0:
        return 0.0
    cov = sum((a - mean_a) * (b - mean_b) for a, b in zip(values_a, values_b))
    return round(cov / (var_a * var_b) ** 0.5, 3)


def compare(shot_metrics, reference_metrics):
    """Delta between a render and a reference, plus a 0..1 similarity score."""
    luma = shot_metrics["mean_luma"] - reference_metrics["mean_luma"]
    contrast_ratio = (shot_metrics["rms_contrast"] /
                      reference_metrics["rms_contrast"]) if reference_metrics["rms_contrast"] else 1.0
    edge_ratio = (shot_metrics["edge_density"] /
                  reference_metrics["edge_density"]) if reference_metrics["edge_density"] else 1.0
    palette = _palette_distance(shot_metrics["dominant"], reference_metrics["dominant"])
    composition = _correlation(shot_metrics["grid_luma"], reference_metrics["grid_luma"])

    # Score: luminance within 25, contrast and edge ratios near 1, palette near 0,
    # composition correlated. Each term is clamped to 0..1 and equally weighted.
    def near(value, target, span):
        return max(0.0, 1.0 - abs(value - target) / span)

    score = (
        near(luma, 0.0, 25.0)
        + near(contrast_ratio, 1.0, 1.0)
        + near(edge_ratio, 1.0, 1.5)
        + (1.0 if palette is None else near(palette, 0.0, 120.0))
        + max(0.0, composition)
    ) / 5.0

    return {
        "delta_luma": round(luma, 1),
        "contrast_ratio": round(contrast_ratio, 2),
        "edge_ratio": round(edge_ratio, 2),
        "palette_distance": palette,
        "composition_correlation": composition,
        "score": round(score, 3),
    }


# --- feedback ---------------------------------------------------------------
# Each rule maps a measured delta to an action phrased for the Blender skills.


def feedback(name, shot, reference, delta):
    notes = []
    if delta["edge_ratio"] < 0.6:
        notes.append(
            f"**Relief and silhouette detail** are {1 - delta['edge_ratio']:.0%} below "
            f"the reference (`blender-modeler` / `archviz`): add window rhythm, "
            f"pilasters, cornices or quoins that cast a shadow line.")
    elif delta["edge_ratio"] > 1.8:
        notes.append(
            "**Too busy** against the reference: the render has far more edge detail. "
            "Check for stray kit pieces, z-fighting or duplicated geometry.")

    if delta["delta_luma"] <= -15:
        notes.append(
            f"**Too dark** by {abs(delta['delta_luma']):.0f} luma (`materials` / "
            f"`lighting`): raise the material albedo or the exposure.")
    elif delta["delta_luma"] >= 15:
        notes.append(
            f"**Too bright** by {delta['delta_luma']:.0f} luma (`materials` / "
            f"`lighting`): lower the albedo or the exposure.")

    if delta["contrast_ratio"] < 0.7:
        notes.append(
            "**Value range compressed** (`lookdev`): the reference has much stronger "
            "light/shadow separation. Deepen recesses or widen the material range.")
    elif delta["contrast_ratio"] > 1.5:
        notes.append(
            "**Value range too harsh** (`lookdev`): soften the key or the shadow "
            "contrast so the surface does not clip.")

    if delta["palette_distance"] is not None and delta["palette_distance"] > 60:
        notes.append(
            f"**Hue drift** ({delta['palette_distance']:.0f} palette distance, "
            f"`materials`): compare the dominant colours in the iteration report "
            f"against the reference and correct the base colours.")

    if delta["composition_correlation"] < 0.3:
        notes.append(
            "**Framing differs** (`camera-cinematography`): the 3x3 luma grid does "
            "not match the reference. Check the camera height, lens and how much "
            "sky/ground is in frame before changing the model.")

    if not notes:
        notes.append(
            "No metric is outside its tolerance. Differences left are stylistic; "
            "confirm visually rather than by number.")
    return notes


# --- verification -----------------------------------------------------------


def verify(targets, judge=False):
    """Check that each shot exists, is not broken, and looks like its expectation."""
    results = []
    for target in targets:
        shot = os.path.join(ROOT, target["shot"])
        entry = {"name": target["name"], "shot": target["shot"], "expect": target["expect"]}
        if not os.path.exists(shot):
            entry["status"] = "MISSING"
            results.append(entry)
            continue
        m = metrics(shot)
        problems = []
        if m["edge_density"] < MIN_EDGE_DENSITY:
            problems.append(f"almost flat (edge density {m['edge_density']})")
        if m["rms_contrast"] < MIN_CONTRAST:
            problems.append(f"no tonal range (contrast {m['rms_contrast']})")
        if m["mean_luma"] < MIN_MEAN_LUMA:
            problems.append(f"black frame (luma {m['mean_luma']})")
        expected_size = target.get("size")
        if expected_size and m["size"] != expected_size:
            problems.append(f"size {m['size']} != {expected_size}")
        entry["metrics"] = {k: m[k] for k in
                            ("size", "mean_luma", "rms_contrast", "edge_density")}
        entry["status"] = "FAIL" if problems else "OK"
        entry["problems"] = problems

        if judge and os.environ.get("OPENROUTER_API_KEY") or (
                judge and os.path.exists(os.path.join(ROOT, ".env"))):
            entry["judge"] = run_judge(shot, target["expect"])
        results.append(entry)
    return results


def run_judge(shot, expect):
    """Advisory vision check. Never a gate: the model is unreliable on this content."""
    import subprocess
    try:
        proc = subprocess.run(
            [sys.executable, os.path.join(ROOT, "Tools", "judge_screenshot.py"),
             "--image", shot, "--expect", expect],
            capture_output=True, text=True, timeout=180, cwd=ROOT)
        out = (proc.stdout or "") + (proc.stderr or "")
        for line in out.splitlines():
            if line.startswith("verdict="):
                return line.strip()
        return "unavailable"
    except Exception as exc:  # noqa: BLE001 - advisory only
        return f"error: {exc}"


# --- reporting --------------------------------------------------------------


def write_iteration(iteration, rows, args):
    """Write the iteration artefacts (gitignored) and return the directory."""
    out_dir = os.path.join(ITER_ROOT, str(iteration))
    os.makedirs(out_dir, exist_ok=True)
    for row in rows:
        shot = os.path.join(ROOT, row["shot"])
        if os.path.exists(shot):
            shutil.copy2(shot, os.path.join(out_dir, os.path.basename(shot)))
    with open(os.path.join(out_dir, "metrics.json"), "w", encoding="utf-8") as handle:
        json.dump({"iteration": iteration, "targets": rows}, handle, indent=2)

    lines = [f"# Iteration {iteration} — reference vs render", ""]
    for row in rows:
        lines.append(f"## {row['name']}")
        lines.append("")
        lines.append(f"- shot: `{row['shot']}`")
        if row.get("reference"):
            lines.append(f"- reference: `{row['reference']}`")
        delta = row.get("delta")
        if delta:
            lines.append(
                f"- **score {delta['score']:.2f}** · luma {delta['delta_luma']:+.0f} · "
                f"contrast ×{delta['contrast_ratio']} · edges ×{delta['edge_ratio']} · "
                f"palette {delta['palette_distance']} · composition r={delta['composition_correlation']}")
        lines.append("")
        for note in row.get("feedback", []):
            lines.append(f"- {note}")
        lines.append("")
    path = os.path.join(out_dir, "feedback.md")
    with open(path, "w", encoding="utf-8") as handle:
        handle.write("\n".join(lines) + "\n")
    return out_dir, path


def load_targets():
    with open(TARGETS, encoding="utf-8") as handle:
        return json.load(handle)["targets"]


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--verify", action="store_true",
                    help="gate the current shots against their expectations")
    ap.add_argument("--compare", action="store_true",
                    help="measure each shot against its reference(s)")
    ap.add_argument("--feedback", action="store_true", help="include action notes (default with --compare)")
    ap.add_argument("--judge", action="store_true", help="add the advisory vision verdict")
    ap.add_argument("--iteration", default=None, help="iteration label for the output directory")
    args = ap.parse_args(argv)

    if not (args.verify or args.compare):
        args.verify = args.compare = True

    targets = load_targets()

    if args.verify:
        print("== verification ==")
        results = verify(targets, judge=args.judge)
        failed = [r for r in results if r["status"] != "OK"]
        for r in results:
            extra = f" ({'; '.join(r['problems'])})" if r.get("problems") else ""
            judge = f"  [judge: {r['judge']}]" if r.get("judge") else ""
            print(f"  {r['status']:7} {r['name']:32}{extra}{judge}")
        print(f"  -> {len(results) - len(failed)}/{len(results)} shots OK")

    if args.compare:
        print("\n== comparison ==")
        rows = []
        for target in targets:
            shot_path = os.path.join(ROOT, target["shot"])
            if not os.path.exists(shot_path):
                continue
            shot_m = metrics(shot_path)
            best = None
            for ref_rel in target.get("references", []):
                ref_path = os.path.join(ROOT, ref_rel)
                if not os.path.exists(ref_path):
                    continue
                ref_m = metrics(ref_path)
                delta = compare(shot_m, ref_m)
                if best is None or delta["score"] > best[1]["score"]:
                    best = (ref_rel, delta, ref_m)
            if best is None:
                continue
            ref_rel, delta, ref_m = best
            rows.append({
                "name": target["name"],
                "shot": target["shot"],
                "reference": ref_rel,
                "delta": delta,
                "feedback": feedback(target["name"], shot_m, ref_m, delta),
                "expect": target["expect"],
            })
            print(f"  {target['name']:32} score {delta['score']:.2f}  "
                  f"luma {delta['delta_luma']:+.0f}  edges ×{delta['edge_ratio']}")
        if args.iteration is not None:
            out_dir, feedback_path = write_iteration(args.iteration, rows, args)
            print(f"\n  iteration written to {os.path.relpath(out_dir, ROOT)}")
            print(f"  feedback: {os.path.relpath(feedback_path, ROOT)}")
            print("  (gitignored: promote a shot into Screenshots/ only after owner confirmation)")
    return 0


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