"""
verify_frontend_sim.py — scratch harness: simulate the deployed June-9 frontend
(plain wrap, full-res bilinear, no mips) vs the current dev frontend (wrap-mode
detection, mip pyramid, trilinear with per-pixel footprint LOD) on a real
new-geometry bundle.  Answers whether the distance-noise issue (#9) is already
covered by the 2b + P0 frontend code that has not yet reached production.

Usage:
    python verify_frontend_sim.py /tmp/hf_new_bundle.json <tile.jpg> <out_prefix>
"""

import base64
import io
import json
import sys

import numpy as np
from PIL import Image


def positive_modulo(v, d):
    return ((v % d) + d) % d


def mirror_repeat(v):
    f = positive_modulo(v, 2.0)
    return np.where(f <= 1.0, f, 2.0 - f)


def detect_wrap_mode(tex):
    h, w, _ = tex.shape
    if w < 4 or h < 4:
        return "mirror"
    col = lambda a, b: np.mean(np.abs(tex[:, a, :3].astype(float) - tex[:, b, :3].astype(float)))
    row = lambda a, b: np.mean(np.abs(tex[a, :, :3].astype(float) - tex[b, :, :3].astype(float)))
    internal = max((col(w // 2, w // 2 + 1) + row(h // 2, h // 2 + 1)) / 2, 1.5)
    return "wrap" if col(w - 1, 0) < internal * 3 and row(h - 1, 0) < internal * 3 else "mirror"


def build_mips(tex):
    mips = [tex.astype(np.float32)]
    cur = mips[0]
    while cur.shape[0] > 1 or cur.shape[1] > 1:
        h, w, _ = cur.shape
        nh, nw = max(1, h // 2), max(1, w // 2)
        ys = np.minimum(2 * np.arange(nh), h - 1)
        xs = np.minimum(2 * np.arange(nw), w - 1)
        y1 = np.minimum(ys + 1, h - 1)
        x1 = np.minimum(xs + 1, w - 1)
        cur = (cur[np.ix_(ys, xs)] + cur[np.ix_(ys, x1)] + cur[np.ix_(y1, xs)] + cur[np.ix_(y1, x1)]) / 4
        mips.append(cur)
    return mips


def sample_bilinear(tex, x, y, wrap):
    h, w, _ = tex.shape
    x0 = np.clip(np.floor(x).astype(int), 0, w - 1)
    y0 = np.clip(np.floor(y).astype(int), 0, h - 1)
    if wrap:
        x1, y1 = (x0 + 1) % w, (y0 + 1) % h
    else:
        x1, y1 = np.minimum(x0 + 1, w - 1), np.minimum(y0 + 1, h - 1)
    fx = (x - np.floor(x))[:, None]
    fy = (y - np.floor(y))[:, None]
    return (tex[y0, x0] * (1 - fx) * (1 - fy) + tex[y0, x1] * fx * (1 - fy)
            + tex[y1, x0] * (1 - fx) * fy + tex[y1, x1] * fx * fy)


def render(bundle, tex, mode):
    """mode: 'old' = June-9 prod (wrap modulo, full-res bilinear, min(W,H) scale)
             'new' = dev      (wrap detection, mips+trilinear, planeW scale)"""
    w, h = bundle["width"], bundle["height"]
    img = np.array(Image.open(io.BytesIO(base64.b64decode(bundle["pixels"]))).convert("RGB"))
    seg = bundle["segments"][0]
    idx = np.frombuffer(base64.b64decode(seg["mask"]), dtype=np.uint32)
    H = np.array(seg["homography"], dtype=np.float64).reshape(3, 3)
    p = seg["plane"]
    texH, texW = tex.shape[:2]

    ys, xs = idx // w, idx % w
    pw, ph = p["width"], p["height"]
    cx, cy = p["x"] + pw / 2, p["y"] + ph / 2

    if mode == "old":
        repeat_w = max(48.0, min(pw, ph) * 0.22)
    else:
        repeat_w = max(32.0, pw * 0.18)
    repeat_h = repeat_w * (texH / texW)

    def to_plane(px, py):
        pts = np.column_stack([px, py, np.ones(len(px))]) @ H.T
        return pts[:, 0] / pts[:, 2], pts[:, 1] / pts[:, 2]

    fx, fy = to_plane(xs.astype(float), ys.astype(float))
    rx, ry = fx - cx, fy - cy

    if mode == "old":
        u = positive_modulo(rx / repeat_w, 1.0)
        v = positive_modulo(ry / repeat_h, 1.0)
        sample = sample_bilinear(tex.astype(np.float32), u * texW, v * texH, wrap=True)
    else:
        wrap = detect_wrap_mode(tex) == "wrap"
        u = positive_modulo(rx / repeat_w, 1.0) if wrap else mirror_repeat(rx / repeat_w)
        v = positive_modulo(ry / repeat_h, 1.0) if wrap else mirror_repeat(ry / repeat_h)

        fx2, fy2 = to_plane(xs.astype(float) + 1, ys.astype(float))
        fx3, fy3 = to_plane(xs.astype(float), ys.astype(float) + 1)
        tcx, tcy = (rx / repeat_w) * texW, (ry / repeat_h) * texH
        txx = ((fx2 - cx) / repeat_w) * texW
        txy = ((fy2 - cy) / repeat_h) * texH
        tyx = ((fx3 - cx) / repeat_w) * texW
        tyy = ((fy3 - cy) / repeat_h) * texH
        du = np.hypot(txx - tcx, txy - tcy)
        dv = np.hypot(tyx - tcx, tyy - tcy)
        footprint = np.maximum(np.maximum(du, dv), 1e-3)
        lod = np.log2(footprint) + 0.5

        mips = build_mips(tex)
        max_l = len(mips) - 1
        l0 = np.clip(np.floor(lod).astype(int), 0, max_l)
        frac = np.clip(lod - l0, 0, 1)
        sample = np.zeros((len(u), 3), np.float32)
        for level in np.unique(l0):
            m = l0 == level
            a = mips[level]
            sa = sample_bilinear(a, u[m] * a.shape[1], v[m] * a.shape[0], wrap)
            if level < max_l:
                b = mips[level + 1]
                sb = sample_bilinear(b, u[m] * b.shape[1], v[m] * b.shape[0], wrap)
                sa = sa + (sb - sa) * frac[m][:, None]
            sample[m] = sa[:, :3]

    # B4 shade decode (same in both versions)
    sr = seg.get("shadeRange") or [0.55, 1.35]
    shade_map = seg.get("shadeMap")
    if shade_map:
        sm = np.frombuffer(base64.b64decode(shade_map), dtype=np.uint8)
        shade = sr[0] + (sm[idx].astype(np.float32) / 255.0) * (sr[1] - sr[0])
    else:
        shade = np.ones(len(idx), np.float32)

    out = img.copy()
    out[ys, xs] = np.clip(sample[:, :3] * shade[:, None], 0, 255).astype(np.uint8)
    return out


def main():
    bundle_path, tile_path, prefix = sys.argv[1], sys.argv[2], sys.argv[3]
    with open(bundle_path) as f:
        bundle = json.load(f)
    tex = np.array(Image.open(tile_path).convert("RGB"))
    print(f"tile: {tile_path} {tex.shape}  wrapMode={detect_wrap_mode(tex)}")

    old = render(bundle, tex, "old")
    new = render(bundle, tex, "new")
    Image.fromarray(old).save(f"verify_out/{prefix}_old_frontend.png")
    Image.fromarray(new).save(f"verify_out/{prefix}_new_frontend.png")
    side = np.hstack([old, new])
    Image.fromarray(side).save(f"verify_out/{prefix}_frontend_compare.png")
    print(f"saved verify_out/{prefix}_frontend_compare.png  (left=June-9 prod, right=dev)")


if __name__ == "__main__":
    main()