"""
verify_p1_sim.py — scratch harness for P1-1 (fold-breaking) and P1-2
(highlight soft-clip).  Renders a bundle with:
  left  = current production: mirror-repeat + hard clamp
  right = P1: makeSeamless 4-tap blend + per-cell random flips + soft clip

Usage:
    python verify_p1_sim.py /tmp/hf_new_bundle.json <tile> <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 make_seamless(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 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(tex, x, y, wrap=True):
    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]
    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):
    """mode 'mirror' = current prod; 'p1' = seamless + flips + softclip"""
    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"]

    tex = make_seamless(tex_raw) if mode == "p1" else 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

    if mode == "mirror":
        u = mirror_repeat(cu)
        v = mirror_repeat(cv)
        sample = sample_bilinear(tex, u * texW, v * texH, wrap=False)
    else:
        ci = np.floor(cu).astype(np.int64)
        cj = np.floor(cv).astype(np.int64)
        hsh = ((ci * 73856093) ^ (cj * 19349663)).astype(np.uint32)
        u = cu - ci
        v = cv - cj
        u = np.where(hsh & 1, 1 - u, u)
        v = np.where(hsh & 2, 1 - v, v)
        sample = sample_bilinear(tex, u * texW, v * texH, wrap=True)

    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 = sample * shade[:, None]
    lit = soft_clip(lit) if mode == "p1" else np.clip(lit, 0, 255)

    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, "mirror")
    b = render(bundle, tex, "p1")
    Image.fromarray(np.hstack([a, b])).save(f"verify_out/{prefix}_p1_compare.png")
    print(f"saved verify_out/{prefix}_p1_compare.png  (left=mirror/prod, right=P1)")


if __name__ == "__main__":
    main()