"""
verify_p1v2_sim.py — P1-1 v2 certification against the deployed P1-1.

deployed (left):  global 4-tap seamless blend + per-cell random flips.
                  Field result: ghost medallions on stone, mushy wood grain,
                  mirror joint lines on planks.
v2 (right):       masked-shift wrap — interior byte-identical to the source,
                  only a ~12% border band crossfades; plain repeat, no flips.

Usage:
    python verify_p1v2_sim.py <bundle.json> <tile> <out_prefix>
"""

import base64
import io
import json
import sys

import numpy as np
from PIL import Image


def make_seamless_4tap(tex):
    h, w, _ = tex.shape
    half_w, half_h = w // 2, h // 2

    def win(n):
        t = 1 - np.abs(2 * np.arange(n) / (n - 1) - 1)
        return t * t * (3 - 2 * t)

    wx = win(w)[None, :, None]
    wy = win(h)[:, None, None]
    t0 = tex.astype(np.float32)
    t1 = np.roll(t0, -half_w, axis=1)
    t2 = np.roll(t0, -half_h, axis=0)
    t3 = np.roll(t1, -half_h, axis=0)
    out = t0 * (wx * wy) + t1 * ((1 - wx) * wy) + t2 * (wx * (1 - wy)) + t3 * ((1 - wx) * (1 - wy))
    return np.clip(out, 0, 255).astype(np.uint8)


def make_wrappable_masked_shift(tex):
    h, w, _ = tex.shape
    band_x = max(2, round(w * 0.12))
    band_y = max(2, round(h * 0.12))

    def edge(n, band):
        i = np.arange(n)
        t = np.minimum(np.minimum(i, n - 1 - i) / band, 1.0)
        return t * t * (3 - 2 * t)

    m = (edge(w, band_x)[None, :] * edge(h, band_y)[:, None])[..., None]
    shifted = np.roll(np.roll(tex, -(w // 2), axis=1), -(h // 2), axis=0).astype(np.float32)
    out = tex.astype(np.float32) * m + shifted * (1 - m)
    return np.clip(out, 0, 255).astype(np.uint8)


def soft_clip(v, knee=220.0):
    rng = 255.0 - knee
    t = np.maximum(v - knee, 0.0)
    return np.where(v <= knee, v, knee + t * rng / (t + rng))


def sample_bilinear_wrap(tex, x, y):
    h, w, _ = tex.shape
    x0 = np.floor(x).astype(int) % w
    y0 = np.floor(y).astype(int) % h
    x1, y1 = (x0 + 1) % w, (y0 + 1) % h
    fx = (x - np.floor(x))[:, None]
    fy = (y - np.floor(y))[:, None]
    t = tex.astype(np.float32)
    return (t[y0, x0] * (1 - fx) * (1 - fy) + t[y0, x1] * fx * (1 - fy)
            + t[y1, x0] * (1 - fx) * fy + t[y1, x1] * fx * fy)


def render(bundle, tex_raw, mode):
    w, h = bundle["width"], bundle["height"]
    raw = base64.b64decode(bundle["pixels"])
    if len(raw) == w * h * 4:
        img = np.frombuffer(raw, np.uint8).reshape(h, w, 4)[:, :, :3].copy()
    else:
        img = np.array(Image.open(io.BytesIO(raw)).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"]

    tex = make_seamless_4tap(tex_raw) if mode == "p1" else make_wrappable_masked_shift(tex_raw)
    texH, texW = tex.shape[:2]

    ys, xs = idx // w, idx % w
    pw = p["width"]
    cx, cy = p["x"] + pw / 2, p["y"] + p["height"] / 2
    repeat_w = max(32.0, pw * 0.18)
    repeat_h = repeat_w * (texH / texW)

    pts = np.column_stack([xs, ys, np.ones(len(xs))]) @ H.T
    fx = pts[:, 0] / pts[:, 2]
    fy = pts[:, 1] / pts[:, 2]
    cu = (fx - cx) / repeat_w
    cv = (fy - cy) / repeat_h

    u = cu - np.floor(cu)
    v = cv - np.floor(cv)
    if mode == "p1":  # deployed: per-cell random flips
        ci = np.floor(cu).astype(np.int64)
        cj = np.floor(cv).astype(np.int64)
        hsh = ((ci * 73856093) ^ (cj * 19349663)).astype(np.uint32)
        u = np.where(hsh & 1, 1 - u, u)
        v = np.where(hsh & 2, 1 - v, v)
    sample = sample_bilinear_wrap(tex, u * texW, v * texH)

    sr = seg.get("shadeRange") or [0.55, 1.35]
    sm = np.frombuffer(base64.b64decode(seg["shadeMap"]), dtype=np.uint8)
    shade = sr[0] + (sm[idx].astype(np.float32) / 255.0) * (sr[1] - sr[0])

    lit = soft_clip(sample * shade[:, None])
    out = img.copy()
    out[ys, xs] = np.clip(lit, 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}")

    a = render(bundle, tex, "p1")
    b = render(bundle, tex, "v2")
    Image.fromarray(np.hstack([a, b])).save(f"verify_out/{prefix}_v2_compare.png")
    print(f"saved verify_out/{prefix}_v2_compare.png  (left=deployed P1, right=v2)")


if __name__ == "__main__":
    main()