"""R4-1 — intrinsic vs heuristic shading evaluation on the reference bundles.

LOCAL evaluation tool (needs torch + the intrinsic package + model weights);
not part of `make verify` — the CI-safe integration harness is
verify_r4_intrinsic_sim.py.

For each committed reference bundle this runs the REAL build_shade_map
(heuristic) and build_intrinsic_shade_map (intrinsic) from app.py on the
bundle's photo + floor mask, then writes a panel to verify_out/:

    original | heuristic shade applied to flat gray | intrinsic shade applied

Applying the decoded shade to a flat gray floor is the most direct artifact
view: any tile pattern, stain or halo visible in the gray region is shading
transfer that would contaminate every replacement floor.

Usage:  python verify_r4_eval.py
"""

import base64
import io
import json
import os

import cv2
import numpy as np
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(HERE, "verify_out")

BUNDLES = [
    ("desk", os.path.join(HERE, "data", "current_bundle.vizbundle.json")),
    ("kitchen", os.path.join(HERE, "data", "ref_kitchen.vizbundle.json")),
]

# --- real implementations from app.py ----------------------------------------
src = open(os.path.join(HERE, "app.py")).read()
ns = {"np": np, "cv2": cv2}
for fn in [
    "_adaptive_shade_range",
    "_encode_shade",
    "_dominant_period",
    "_suppress_periodic_shading",
    "build_shade_map",
    "build_intrinsic_shade_map",
]:
    start = src.index(f"def {fn}")
    end = src.index("\ndef ", start + 10)
    exec(compile(src[start:end], "app.py", "exec"), ns)


def load_bundle(path):
    d = json.load(open(path))
    img = np.asarray(
        Image.open(io.BytesIO(base64.b64decode(d["pixels"]))).convert("RGB")
    )
    h, w = d["height"], d["width"]
    mask = np.zeros(w * h, bool)
    for s in d["segments"]:
        idx = np.frombuffer(base64.b64decode(s["mask"]), dtype=np.uint32)
        mask[idx] = True
    return img, mask.reshape(h, w).astype(np.uint8)


def decode(enc, rng):
    lo, hi = rng
    return lo + enc.astype(np.float64) / 255.0 * (hi - lo)


def shade_on_gray(img, mask, rel):
    """Composite: flat gray floor x shade over the original photo."""
    out = img.astype(np.float64).copy()
    gray = 205.0 * np.clip(rel, 0.0, 2.0)
    for c in range(3):
        ch = out[:, :, c]
        ch[mask > 0] = np.clip(gray[mask > 0], 0, 255)
    return out.astype(np.uint8)


def main():
    os.makedirs(OUT, exist_ok=True)

    print("loading intrinsic model (v2)...", flush=True)
    # same headless-trust shim as app._load_intrinsic_model
    import torch.hub as _hub

    os.makedirs(_hub.get_dir(), exist_ok=True)
    tl = os.path.join(_hub.get_dir(), "trusted_list")
    if "rwightman_gen-efficientnet-pytorch" not in (
        open(tl).read() if os.path.exists(tl) else ""
    ):
        with open(tl, "a") as f:
            f.write("rwightman_gen-efficientnet-pytorch\n")

    from intrinsic.pipeline import load_models

    ns["device"] = "cpu"
    ns["intrinsic_models"] = load_models("v2", device="cpu")
    print("model loaded.", flush=True)

    for name, path in BUNDLES:
        img, mask = load_bundle(path)
        h, w = mask.shape

        enc_h, rng_h = ns["build_shade_map"](img, mask)
        import time

        t0 = time.perf_counter()
        enc_i, rng_i = ns["build_intrinsic_shade_map"](img, mask)
        dt = time.perf_counter() - t0

        if enc_h is None or enc_i is None:
            print(f"  [{name}] FAILED: heuristic={enc_h is not None} intrinsic={enc_i is not None}")
            continue

        rel_h = decode(enc_h.reshape(h, w), rng_h)
        rel_i = decode(enc_i.reshape(h, w), rng_i)

        for label, rel in (("heuristic", rel_h), ("intrinsic", rel_i)):
            v = rel[mask > 0]
            print(
                f"  [{name}] {label:9s} p5={np.percentile(v,5):.3f} "
                f"p50={np.percentile(v,50):.3f} p95={np.percentile(v,95):.3f} "
                f"range=({min(rng_h if label=='heuristic' else rng_i):.2f},"
                f"{max(rng_h if label=='heuristic' else rng_i):.2f})"
            )
        print(f"  [{name}] intrinsic runtime: {dt:.1f}s on cpu", flush=True)

        panel = np.concatenate(
            [img, shade_on_gray(img, mask, rel_h), shade_on_gray(img, mask, rel_i)],
            axis=1,
        )
        scale = min(2200 / panel.shape[1], 1.0)
        if scale < 1.0:
            panel = cv2.resize(
                panel, (round(panel.shape[1] * scale), round(panel.shape[0] * scale))
            )
        out_path = os.path.join(OUT, f"r4_eval_{name}.png")
        Image.fromarray(panel).save(out_path)
        print(f"  [{name}] panel: {out_path}  (original | heuristic | intrinsic)")

    return 0


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