verify.py

"""
verify.py — certify B1/B4/B6/B8 from a .vizbundle.json bundle file.

Usage:
    python verify.py data/current_bundle.vizbundle.json [--out verify_out]
"""

import argparse
import base64
import json
import os
import sys

import numpy as np
from PIL import Image, ImageDraw


# ---------------------------------------------------------------------------
# helpers
# ---------------------------------------------------------------------------

def decode_shade_map(b64: str, width: int, height: int) -> np.ndarray:
    raw = base64.b64decode(b64)
    arr = np.frombuffer(raw, dtype=np.uint8)
    if arr.size != width * height:
        raise ValueError(
            f"shadeMap size mismatch: got {arr.size}, expected {width * height}"
        )
    return arr.reshape(height, width)


def quad_points(flat: list[float]) -> list[tuple[int, int]]:
    """Convert flat [x0,y0,x1,y1,...] list to (x,y) tuples."""
    pts = [(int(flat[i]), int(flat[i + 1])) for i in range(0, len(flat), 2)]
    # close polygon
    return pts + [pts[0]]


# ---------------------------------------------------------------------------
# checks
# ---------------------------------------------------------------------------

def check_b1_shade_map(seg: dict, width: int, height: int, out_dir: str) -> bool:
    print("\n[B1] shadeMap")
    raw = seg.get("shadeMap")
    if not raw:
        print("  FAIL — shadeMap missing")
        return False

    shade = decode_shade_map(raw, width, height)
    mn, mx, mean = int(shade.min()), int(shade.max()), float(shade.mean())
    print(f"  shape : {shade.shape}")
    print(f"  range : [{mn}, {mx}]  mean={mean:.1f}")

    # save greyscale visualisation
    img = Image.fromarray(shade, mode="L")
    path = os.path.join(out_dir, "b1_shade_map.png")
    img.save(path)
    print(f"  saved : {path}")

    if mn == mx:
        print("  WARN  — shade map is flat (all one value)")
    print("  OK")
    return True


def check_b4_shade_range(seg: dict) -> bool:
    print("\n[B4] shadeRange")
    sr = seg.get("shadeRange")
    if sr is None:
        print("  FAIL — shadeRange missing")
        return False
    lo, hi = sr
    print(f"  lo={lo:.4f}  hi={hi:.4f}  span={hi - lo:.4f}")
    if hi <= lo:
        print("  FAIL — hi must be > lo")
        return False
    if lo < 0 or hi > 4:
        print(f"  WARN  — range [{lo:.3f}, {hi:.3f}] looks unusual (expected ~0.5–2.5)")
    print("  OK")
    return True


def check_b6_vanishing_points(seg: dict, width: int, height: int, out_dir: str) -> bool:
    print("\n[B6] vanishingPoints")
    plane = seg.get("plane", {})
    vp1 = plane.get("vanishingPoint")
    vp2 = plane.get("vanishingPoint2")

    if not vp1:
        print("  FAIL — vanishingPoint missing from plane")
        return False

    print(f"  VP1 : ({vp1['x']:.1f}, {vp1['y']:.1f})")
    vp2_str = f"({vp2['x']:.1f}, {vp2['y']:.1f})" if vp2 else "null (single-VP room)"
    print(f"  VP2 : {vp2_str}")

    # draw on a blank canvas
    img = Image.new("RGB", (width, height), (30, 30, 30))
    draw = ImageDraw.Draw(img)
    r = max(8, width // 80)
    x1, y1 = int(vp1["x"]), int(vp1["y"])
    draw.ellipse([x1 - r, y1 - r, x1 + r, y1 + r], fill=(255, 80, 80), outline=(255, 255, 255))
    draw.text((x1 + r + 4, y1 - r), "VP1", fill=(255, 80, 80))
    if vp2:
        x2, y2 = int(vp2["x"]), int(vp2["y"])
        draw.ellipse([x2 - r, y2 - r, x2 + r, y2 + r], fill=(80, 180, 255), outline=(255, 255, 255))
        draw.text((x2 + r + 4, y2 - r), "VP2", fill=(80, 180, 255))

    path = os.path.join(out_dir, "b6_vanishing_points.png")
    img.save(path)
    print(f"  saved : {path}")
    print("  OK")
    return True


def check_b8_quad_vs_hull(seg: dict, width: int, height: int, out_dir: str) -> bool:
    print("\n[B8] quad vs hullQuad")
    plane = seg.get("plane", {})
    quad_flat = plane.get("quad")
    hull_flat = plane.get("hullQuad")

    if not quad_flat:
        print("  FAIL — quad missing from plane")
        return False
    if not hull_flat:
        print("  FAIL — hullQuad missing from plane")
        return False

    quad_pts = quad_points(quad_flat)
    hull_pts = quad_points(hull_flat)
    print(f"  quad     : {quad_pts[:-1]}")
    print(f"  hullQuad : {hull_pts[:-1]}")

    img = Image.new("RGB", (width, height), (20, 20, 20))
    draw = ImageDraw.Draw(img)
    draw.line(hull_pts, fill=(255, 200, 0), width=3)
    draw.line(quad_pts, fill=(0, 220, 100), width=2)

    # legend
    draw.rectangle([10, 10, 26, 20], fill=(255, 200, 0))
    draw.text((30, 10), "hullQuad", fill=(255, 200, 0))
    draw.rectangle([10, 26, 26, 36], fill=(0, 220, 100))
    draw.text((30, 26), "quad (fitted)", fill=(0, 220, 100))

    path = os.path.join(out_dir, "b8_quad_overlay.png")
    img.save(path)
    print(f"  saved : {path}")
    print("  OK")
    return True


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(description="Verify a vizbundle JSON file.")
    parser.add_argument("bundle", help="Path to .vizbundle.json")
    parser.add_argument("--out", default="verify_out", help="Output directory (default: verify_out)")
    args = parser.parse_args()

    with open(args.bundle) as f:
        bundle = json.load(f)

    width = bundle["width"]
    height = bundle["height"]
    segments = bundle.get("segments", [])

    if not segments:
        print("ERROR — no segments in bundle")
        sys.exit(1)

    os.makedirs(args.out, exist_ok=True)
    print(f"Bundle: {args.bundle}  ({width}x{height}, {len(segments)} segment(s))")

    results = []
    for i, seg in enumerate(segments):
        print(f"\n=== Segment {i} — {seg.get('className', '?')} ===")
        results.append(check_b1_shade_map(seg, width, height, args.out))
        results.append(check_b4_shade_range(seg))
        results.append(check_b6_vanishing_points(seg, width, height, args.out))
        results.append(check_b8_quad_vs_hull(seg, width, height, args.out))

    print("\n" + "=" * 40)
    passed = sum(results)
    total = len(results)
    print(f"RESULT: {passed}/{total} checks passed")
    if passed < total:
        sys.exit(1)


if __name__ == "__main__":
    main()