"""
verify_geometry.py — scratch harness for the T1 perspective fix.

Renders a procedural brick grid through the OLD and NEW homography
constructions on a real bundle, side by side, so the foreshortening
difference is directly visible.  Geometry functions are copied from app.py
(pure cv2/numpy) so this runs without torch/transformers installed.

Usage:
    python verify_geometry.py data/current_bundle.vizbundle.json
"""

import base64
import io
import json
import sys

import cv2
import numpy as np
from PIL import Image


# --- copied from app.py (pure geometry helpers) -----------------------------

def fit_floor_edges(mask):
    h, w = mask.shape[:2]
    row_ys, lefts, rights = [], [], []
    step = max(1, h // 260)
    for y in range(0, h, step):
        row_xs = np.where(mask[y] > 0)[0]
        if len(row_xs) < max(8, w * 0.01):
            continue
        row_ys.append(float(y))
        lefts.append(float(np.percentile(row_xs, 3)))
        rights.append(float(np.percentile(row_xs, 97)))
    if len(row_ys) < 8:
        return None
    row_ys_np = np.asarray(row_ys, dtype=np.float32)
    return np.polyfit(row_ys_np, np.asarray(lefts, dtype=np.float32), 1), np.polyfit(
        row_ys_np, np.asarray(rights, dtype=np.float32), 1
    )


def convex_hull_quad(mask):
    ys, xs = np.where(mask > 0)
    if len(xs) < 50:
        return None
    pts = np.column_stack([xs, ys]).astype(np.float32)
    hull = cv2.convexHull(pts)
    if hull is None or len(hull) < 4:
        return None
    rect = cv2.minAreaRect(hull.squeeze())
    box = cv2.boxPoints(rect)
    h, w = mask.shape[:2]
    box[:, 0] = np.clip(box[:, 0], 0, w - 1)
    box[:, 1] = np.clip(box[:, 1], 0, h - 1)
    return box


def detect_dual_vanishing_points(img_np, floor_mask):
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    edges = cv2.Canny(gray, 60, 160)
    edges[floor_mask == 0] = 0
    lines = cv2.HoughLinesP(
        edges, rho=1, theta=np.pi / 180, threshold=60,
        minLineLength=max(40, min(img_np.shape[:2]) // 16), maxLineGap=24,
    )
    if lines is None:
        return None, None

    h, w = img_np.shape[:2]
    pos_lines, neg_lines = [], []
    for line in lines[:, 0, :]:
        x1, y1, x2, y2 = [float(v) for v in line]
        dx, dy = x2 - x1, y2 - y1
        length = float(np.hypot(dx, dy))
        if length < 40 or abs(dx) < 1:
            continue
        slope = dy / dx
        if abs(slope) < 0.18:
            continue
        entry = (x1, y1, x2, y2, slope, length)
        (pos_lines if slope > 0 else neg_lines).append(entry)

    def _find_vp(group):
        intersections = []
        for i, (x1, y1, _, _, s1, l1) in enumerate(group):
            a1 = y1 - s1 * x1
            for x3, y3, _, _, s2, l2 in group[i + 1:]:
                if abs(s1 - s2) < 0.08:
                    continue
                denom = s1 - s2
                if abs(denom) < 1e-9:
                    continue
                a2 = y3 - s2 * x3
                x = (a2 - a1) / denom
                y = s1 * x + a1
                if -w * 0.6 <= x <= w * 1.6 and -h * 1.2 <= y <= h * 1.0:
                    intersections.append((x, y, min(l1, l2)))
        if len(intersections) < 3:
            return None
        pts = np.array([[p[0], p[1]] for p in intersections], np.float32)
        weights = np.array([p[2] for p in intersections], np.float32)
        center = np.average(pts, axis=0, weights=weights)
        dist = np.linalg.norm(pts - center, axis=1)
        keep = dist <= np.percentile(dist, 70)
        if keep.sum() >= 3:
            center = np.average(pts[keep], axis=0, weights=weights[keep])
        return {"x": float(center[0]), "y": float(center[1])}

    vp_right = _find_vp(pos_lines)
    vp_left = _find_vp(neg_lines)
    candidates = [(vp, abs(vp["y"])) for vp in [vp_right, vp_left] if vp is not None]
    if not candidates:
        return None, None
    candidates.sort(key=lambda t: t[1])
    primary = candidates[0][0]
    secondary = candidates[1][0] if len(candidates) > 1 else None
    return primary, secondary


def _shared_quad_setup(mask):
    """Common src-quad construction shared by old/new paths (pre-VP)."""
    ys, xs = np.where(mask > 0)
    xs_f, ys_f = xs.astype(np.float32), ys.astype(np.float32)
    x1, x2 = float(np.percentile(xs_f, 1)), float(np.percentile(xs_f, 99))
    y1, y2 = float(np.percentile(ys_f, 1)), float(np.percentile(ys_f, 99))
    width, height = x2 - x1, y2 - y1

    top_y = float(np.percentile(ys_f, 8))
    bottom_y = float(np.percentile(ys_f, 97))
    left_fit, right_fit = fit_floor_edges(mask)

    top_left = float(np.polyval(left_fit, top_y))
    top_right = float(np.polyval(right_fit, top_y))
    bottom_left = float(np.polyval(left_fit, bottom_y))
    bottom_right = float(np.polyval(right_fit, bottom_y))
    lower_xs = xs_f[ys_f >= np.percentile(ys_f, 80)]
    bottom_left = min(bottom_left, float(np.percentile(lower_xs, 4)))
    bottom_right = max(bottom_right, float(np.percentile(lower_xs, 96)))

    min_top_width = max(24.0, width * 0.18)
    top_center = (top_left + top_right) * 0.5
    if top_right - top_left < min_top_width:
        top_left = top_center - min_top_width * 0.5
        top_right = top_center + min_top_width * 0.5

    min_bottom_width = max(min_top_width * 1.25, width * 0.45)
    bottom_center = (bottom_left + bottom_right) * 0.5
    if bottom_right - bottom_left < min_bottom_width:
        bottom_left = bottom_center - min_bottom_width * 0.5
        bottom_right = bottom_center + min_bottom_width * 0.5

    h, w = mask.shape[:2]
    src = np.float32([
        [np.clip(bottom_left, 0, w - 1), np.clip(bottom_y, 0, h - 1)],
        [np.clip(bottom_right, 0, w - 1), np.clip(bottom_y, 0, h - 1)],
        [np.clip(top_right, 0, w - 1), np.clip(top_y, 0, h - 1)],
        [np.clip(top_left, 0, w - 1), np.clip(top_y, 0, h - 1)],
    ])
    return src, (x1, x2, y1, y2, width, height, top_y, bottom_y, top_center)


def estimate_old(mask, img_np):
    src, (x1, x2, y1, y2, width, height, top_y, bottom_y, top_center) = _shared_quad_setup(mask)
    h, w = mask.shape[:2]

    vanishing_point, _ = detect_dual_vanishing_points(img_np, mask)
    if vanishing_point is not None and vanishing_point["y"] < bottom_y:
        vp_x = float(np.clip(vanishing_point["x"], -w * 0.25, w * 1.25))
        top_width = max(src[2][0] - src[3][0], width * 0.16)
        horizon_gap = max(bottom_y - top_y, 1.0)
        convergence = np.clip((top_y - vanishing_point["y"]) / horizon_gap, 0.12, 0.75)
        top_center = top_center * (1 - convergence * 0.35) + vp_x * (convergence * 0.35)
        src[3][0] = np.clip(top_center - top_width * 0.5, 0, w - 1)
        src[2][0] = np.clip(top_center + top_width * 0.5, 0, w - 1)

    hull_box = convex_hull_quad(mask)
    if hull_box is not None:
        src[0][0] = min(src[0][0], float(np.min(hull_box[:, 0])))
        src[1][0] = max(src[1][0], float(np.max(hull_box[:, 0])))
        src[0][1] = src[1][1] = max(src[0][1], float(np.max(hull_box[:, 1])))
        src[2][1] = src[3][1] = min(src[2][1], float(np.min(hull_box[:, 1])))
        src = np.clip(src, [0, 0], [w - 1, h - 1]).astype(np.float32)

    dst = np.float32([[x1, y2], [x2, y2], [x2, y1], [x1, y1]])
    H = cv2.getPerspectiveTransform(src, dst)
    return H, {"x": x1, "y": y1, "width": width, "height": height}, src


def estimate_new(mask, img_np):
    src, (x1, x2, y1, y2, width, height, top_y, bottom_y, top_center) = _shared_quad_setup(mask)
    h, w = mask.shape[:2]

    vanishing_point, _ = detect_dual_vanishing_points(img_np, mask)

    hull_box = convex_hull_quad(mask)
    if hull_box is not None:
        src[0][0] = min(src[0][0], float(np.min(hull_box[:, 0])))
        src[1][0] = max(src[1][0], float(np.max(hull_box[:, 0])))
        src[0][1] = src[1][1] = max(src[0][1], float(np.max(hull_box[:, 1])))
        src[2][1] = src[3][1] = min(src[2][1], float(np.min(hull_box[:, 1])))
        src = np.clip(src, [0, 0], [w - 1, h - 1]).astype(np.float32)

    bottom_y_f, top_y_f = float(src[0][1]), float(src[3][1])
    bl_x, br_x = float(src[0][0]), float(src[1][0])
    span_y = max(bottom_y_f - top_y_f, 1.0)
    bottom_w = max(br_x - bl_x, 1.0)

    if vanishing_point is not None and vanishing_point["y"] < top_y_f - 4:
        vp_x = float(np.clip(vanishing_point["x"], -w * 0.5, w * 1.5))
        vp_y = float(vanishing_point["y"])
    else:
        vp_x = (bl_x + br_x) * 0.5
        vp_y = top_y_f - 0.35 * span_y
    vp_y = min(vp_y, top_y_f - 0.08 * span_y)

    t = (top_y_f - bottom_y_f) / (vp_y - bottom_y_f)
    src[3][0] = bl_x + (vp_x - bl_x) * t
    src[2][0] = br_x + (vp_x - br_x) * t
    if src[2][0] - src[3][0] < bottom_w * 0.06:
        top_cx = (float(src[2][0]) + float(src[3][0])) * 0.5
        src[3][0] = top_cx - bottom_w * 0.03
        src[2][0] = top_cx + bottom_w * 0.03

    depth_ratio = (bottom_y_f - vp_y) / max(top_y_f - vp_y, 1e-3)
    dst_h = width * float(w) * (depth_ratio - 1.0) / bottom_w
    dst_h = float(np.clip(dst_h, height * 0.8, height * 5.0))

    dst = np.float32([[x1, y2], [x2, y2], [x2, y2 - dst_h], [x1, y2 - dst_h]])
    H = cv2.getPerspectiveTransform(src, dst)
    print(f"  VP used: ({vp_x:.0f}, {vp_y:.0f})  detected={vanishing_point is not None}")
    print(f"  top width: {src[2][0]-src[3][0]:.0f}px vs bottom {bottom_w:.0f}px "
          f"(ratio {(src[2][0]-src[3][0])/bottom_w:.2f})")
    print(f"  depth_ratio: {depth_ratio:.2f}  dst_h: {dst_h:.0f} (bbox h was {height:.0f})")
    return H, {"x": x1, "y": y2 - dst_h, "width": width, "height": dst_h}, src


# --- renderer ----------------------------------------------------------------

def render_grid(img_np, mask, H, plane, label, repeat_mode):
    """Inverse-map every mask pixel through H and paint a brick grid."""
    out = img_np.copy()
    ys, xs = np.where(mask > 0)
    pts = np.column_stack([xs, ys]).astype(np.float64)
    ones = np.ones((len(pts), 1))
    p = np.hstack([pts, ones]) @ H.T
    fx = p[:, 0] / p[:, 2]
    fy = p[:, 1] / p[:, 2]

    pw, ph = plane["width"], plane["height"]
    if repeat_mode == "old":
        repeat = max(48.0, min(pw, ph) * 0.22)
    else:
        repeat = max(32.0, pw * 0.18)

    u = (fx - plane["x"]) / repeat
    v = (fy - plane["y"]) / repeat

    # brick pattern: offset every other row by half a tile, dark grout lines
    row = np.floor(v).astype(int)
    uu = u + (row % 2) * 0.5
    cell = ((np.floor(uu).astype(int) + row) % 2).astype(bool)
    fu, fv = uu - np.floor(uu), v - np.floor(v)
    grout = (fu < 0.06) | (fv < 0.06)

    color = np.where(cell[:, None], [184, 115, 51], [222, 184, 135]).astype(np.uint8)
    color[grout] = (60, 60, 60)
    out[ys, xs] = (0.75 * color + 0.25 * out[ys, xs]).astype(np.uint8)
    return out


def main():
    bundle_path = sys.argv[1] if len(sys.argv) > 1 else "data/current_bundle.vizbundle.json"
    with open(bundle_path) as f:
        bundle = json.load(f)

    w, h = bundle["width"], bundle["height"]
    img = Image.open(io.BytesIO(base64.b64decode(bundle["pixels"]))).convert("RGB")
    img_np = np.array(img)
    print(f"Image: {img_np.shape}")

    seg = bundle["segments"][0]
    indices = np.frombuffer(base64.b64decode(seg["mask"]), dtype=np.uint32)
    mask = np.zeros(h * w, np.uint8)
    mask[indices] = 1
    mask = mask.reshape(h, w)
    print(f"Mask pixels: {int(mask.sum())}")

    print("\n--- OLD geometry ---")
    H_old, plane_old, src_old = estimate_old(mask, img_np)
    print(f"  src quad: {src_old.flatten().round(0).tolist()}")
    out_old = render_grid(img_np, mask, H_old, plane_old, "old", "old")

    print("\n--- NEW geometry ---")
    H_new, plane_new, src_new = estimate_new(mask, img_np)
    print(f"  src quad: {src_new.flatten().round(0).tolist()}")
    out_new = render_grid(img_np, mask, H_new, plane_new, "new", "new")

    side = np.hstack([out_old, out_new])
    Image.fromarray(out_old).save("verify_out/geometry_old.png")
    Image.fromarray(out_new).save("verify_out/geometry_new.png")
    Image.fromarray(side).save("verify_out/geometry_compare.png")
    print("\nSaved verify_out/geometry_old.png, geometry_new.png, geometry_compare.png")


if __name__ == "__main__":
    main()