"""R0-1 — deterministic golden render: bundle x tile -> PNG.

A faithful Python port of the CURRENT frontend composite (canvas-engine.ts):
texture prep (wrap detection -> period-snap -> masked-shift fallback), mip
pyramid + trilinear with per-pixel footprint LOD, shade-map decode, homography-
mapped light vector + gloss-gated specular, colour cast, soft highlight clip,
confidence-map alpha. Texture-prep and sampling primitives are imported from
verify_n1_sim so this stays in lockstep with the certified implementations.

Usage:
    python golden_render.py <bundle.json[.gz]> <tile-image> <out.png>

The output is resized to max-dim 720 so goldens stay small and stable.
"""

import base64
import gzip
import json
import sys

import numpy as np
from PIL import Image

from verify_n1_sim import (
    build_mips,
    detect_wrap_mode,
    flatten_luminance,
    make_seamless,
    period_snap,
    sample_bilinear_wrap,
)

OUT_MAX_DIM = 720


def load_bundle(path):
    if path.endswith(".gz"):
        with gzip.open(path, "rt") as f:
            return json.load(f)
    return json.load(open(path))


def estimate_gloss(tex):
    """Port of estimateGloss (canvas-engine.ts): mean 4px luminance gradient."""
    lum = tex[:, :, 0] * 0.299 + tex[:, :, 1] * 0.587 + tex[:, :, 2] * 0.114
    a = lum[::4, 4::4]
    b = lum[::4, :-4:4][:, : a.shape[1]]
    mean_grad = float(np.mean(np.abs(a - b))) if a.size else 0.0
    return float(np.clip(1 - mean_grad / 24, 0, 1))


def soft_clip(v):
    """Port of softClipByte: linear below 220, rational shoulder above."""
    knee, rng = 220.0, 35.0
    t = v - knee
    return np.where(v <= knee, v, knee + (t * rng) / (t + rng))


def apply_shade(sample, m):
    """Port of applyShade (canvas-engine.ts) — N4 highlight-preserving shading.

    Dimming (m <= 1) stays a physical multiply. Brightening near-white texels
    by multiply clips at 255 and erases pale-tile detail (the white-tile
    washout), so for m > 1 blend toward a screen-style lift by texel
    brightness: dark texels keep the multiply, bright texels brighten by
    closing their gap to white — grout/figuring contrast survives by
    construction because the result stays below 255.
    """
    mult = sample * m
    screen = 255.0 - (255.0 - sample) / np.maximum(m, 1e-6)
    t = sample / 255.0
    w = t * t * (3.0 - 2.0 * t)
    return np.where(m <= 1.0, mult, mult * (1.0 - w) + screen * w)


def prepare_texture(tile_path):
    tex = np.asarray(Image.open(tile_path).convert("RGB"))
    h, w, _ = tex.shape
    mode, _, _ = detect_wrap_mode(tex)
    repeat_scale = 1.0
    lay_mode = "repeat"
    if mode != "wrap":
        prepared, info = period_snap(tex)
        if info[0] == "snap":
            repeat_scale = prepared.shape[1] / w
            tex = prepared
        else:
            tex = make_seamless(tex)
            # R2-3 v1 — truly organic materials (healed, not periodic) lay as
            # procedural cells; periodic/authored tiles keep the plain repeat.
            lay_mode = "cells"
        # R2-3b — strip the photo's baked lighting (mirrors canvas-engine.ts):
        # snapped patterns stop showing tonal seams at every repeat; healed
        # organic textures stop reading as patchwork across cells.
        tex = flatten_luminance(tex)
    return tex, repeat_scale, lay_mode


_M32 = np.uint64(0xFFFFFFFF)


def hash01(a, b):
    """Port of hash01 (canvas-engine.ts) — bit-exact with Math.imul/>>> JS
    semantics so both engines lay identical cells."""
    a = (np.asarray(a, np.int64) & 0xFFFFFFFF).astype(np.uint64)
    b = (np.asarray(b, np.int64) & 0xFFFFFFFF).astype(np.uint64)
    h = ((a * np.uint64(374761393)) & _M32) ^ ((b * np.uint64(668265263)) & _M32)
    h = ((h ^ (h >> np.uint64(13))) * np.uint64(1274126177)) & _M32
    return ((h ^ (h >> np.uint64(16))) >> np.uint64(8)).astype(np.float64) / 16777216.0


def procedural_cells(rx, ry, repeat_w, repeat_h, grout_frac, tw, th, footprint, pattern="straight"):
    """Port of the R2-3 v1 cell logic (canvas-engine.ts composite):
    per-cell toroidal grain offset, tone jitter, pattern-driven bond
    (straight = random running bond, brick = fixed half offset, grid/diamond
    = aligned), and an anti-aliased seam mask. Returns (u, v, tone, blend)."""
    row = np.floor(ry / repeat_h).astype(np.int64)
    if pattern in ("grid", "diamond"):
        stagger = np.zeros(len(row), np.float64)
    elif pattern == "brick":
        stagger = (row & 1) * 0.5
    else:
        stagger = hash01(row, np.full_like(row, 0x9E37))
    sx = rx + stagger * repeat_w
    col = np.floor(sx / repeat_w).astype(np.int64)
    lu = sx / repeat_w - col
    lv = ry / repeat_h - row
    # R2-3b — grid/diamond tiles are factory prints (real tiles repeat their
    # print): content stays aligned per cell; organic grain keeps the
    # per-cell toroidal window. Mirrors canvas-engine.ts.
    if pattern in ("grid", "diamond"):
        u = np.mod(lu, 1.0)
        v = np.mod(lv, 1.0)
    else:
        u = np.mod(lu + hash01(col, row), 1.0)
        v = np.mod(lv + hash01(col + 0x55, row - 0x21), 1.0)
    tone = 0.94 + 0.12 * hash01(col - 0x13, row + 0x77)
    half_u = grout_frac * 0.5
    half_v = half_u * (repeat_w / repeat_h)
    aa_u = np.maximum(footprint / tw, 1e-4)
    aa_v = np.maximum(footprint / th, 1e-4)
    du_b = np.minimum(lu, 1.0 - lu)
    dv_b = np.minimum(lv, 1.0 - lv)
    b_u = np.clip((du_b - half_u) / aa_u + 0.5, 0.0, 1.0)
    b_v = np.clip((dv_b - half_v) / aa_v + 0.5, 0.0, 1.0)
    return u, v, tone, 1.0 - np.minimum(b_u, b_v)


def apply_reflection(texel, base, xs, ys, h, w):
    """Port of the R4-2 gloss reflection (canvas-engine.ts composite):
    mirror the above-floor scene about the per-column contact line, faded
    with distance and weighted by reflected brightness squared."""
    ys_f = ys.astype(np.float64)
    top = np.full(w, np.inf)
    np.minimum.at(top, xs, ys_f)
    floor_flags = np.zeros((h, w), bool)
    floor_flags[ys, xs] = True
    ty = top[xs]
    ry_refl = np.floor(2.0 * ty - ys_f).astype(np.int64)
    valid = np.isfinite(ty) & (ry_refl >= 0) & (ry_refl < h)
    ry_c = np.clip(ry_refl, 0, h - 1)
    src_floor = floor_flags[ry_c, xs]
    refl = base[ry_c, xs]
    lum = (refl[:, 0] * 0.299 + refl[:, 1] * 0.587 + refl[:, 2] * 0.114) / 255.0
    fade = np.maximum(0.0, 1.0 - (ys_f - ty) / (h * 0.35))
    k = np.where(valid & ~src_floor, 0.22 * fade * lum * lum, 0.0)
    k = np.where(k > 0.003, k, 0.0)
    return texel * (1.0 - k[:, None]) + refl * k[:, None]


def render(bundle_path, tile_path, finish="matte"):
    d = load_bundle(bundle_path)
    w, h = d["width"], d["height"]
    base = np.asarray(
        Image.open(__import__("io").BytesIO(base64.b64decode(d["pixels"]))).convert("RGB")
    ).astype(np.float64)
    seg = max(d["segments"], key=lambda s: len(s["mask"]))

    mask_idx = np.frombuffer(base64.b64decode(seg["mask"]), dtype=np.uint32)
    mask = np.zeros(w * h, bool)
    mask[mask_idx] = True
    mask = mask.reshape(h, w)

    H = np.asarray(seg["homography"], np.float64).reshape(3, 3)
    plane = seg.get("plane") or {}
    plane_w = max(plane.get("width", w), 1)
    plane_h = max(plane.get("height", h), 1)
    plane_cx = plane.get("x", 0) + plane_w / 2
    plane_cy = plane.get("y", 0) + plane_h / 2
    rot_deg = plane.get("defaultRotation") or 0.0
    rad = np.deg2rad(rot_deg)
    cos, sin = np.cos(-rad), np.sin(-rad)

    shade_map = (
        np.frombuffer(base64.b64decode(seg["shadeMap"]), np.uint8).reshape(h, w).astype(np.float64)
        if seg.get("shadeMap")
        else None
    )
    shade_lo, shade_hi = seg.get("shadeRange") or (0.55, 1.35)
    conf = (
        np.frombuffer(base64.b64decode(seg["confidenceMap"]), np.uint8).reshape(h, w).astype(np.float64) / 255.0
        if seg.get("confidenceMap")
        else None
    )
    ct = seg.get("colorTemperature") or {}
    if "cast" in ct:
        ct = ct["cast"]
    col = np.array([ct.get("r", 1.0), ct.get("g", 1.0), ct.get("b", 1.0)])
    lv = seg.get("lightVector")

    tex, repeat_scale, lay_mode = prepare_texture(tile_path)
    # R4-2b — catalog finish overrides the texture-smoothness estimate (T9):
    # matte gets no sheen however smooth the print; gloss keeps a baseline.
    # Mirrors canvas-engine.ts.
    gloss_est = estimate_gloss(np.asarray(Image.open(tile_path).convert("RGB")).astype(np.float64))
    gloss = 0.0 if finish == "matte" else max(0.5, gloss_est) if finish == "gloss" else gloss_est
    # R2-3 v1 — derived seam tone + width; mirrors canvas-engine.ts (same
    # every-16th-pixel mean, same no-physical-metadata fraction).
    grout_col = tex.reshape(-1, 3)[::16].astype(np.float64).mean(axis=0) * 0.45
    grout_frac = 0.012
    th, tw, _ = tex.shape
    mips = build_mips(tex)
    max_l = len(mips) - 1

    # R1-3 — mirror of canvas-engine.ts: metric plane scale when present
    # (pixel-ish or metre plane units alike — backend gates mpu hard),
    # heuristic fallback otherwise (info.scale = 1 in goldens).
    DEFAULT_TILE_M = 0.6
    mpu = plane.get("metersPerUnit")
    repeat_w = 0.0
    if mpu and mpu > 0:
        repeat_w = (DEFAULT_TILE_M / mpu) * repeat_scale
        if not (np.isfinite(repeat_w) and repeat_w > 0):
            repeat_w = 0.0
    if not repeat_w:
        repeat_w = max(48.0, min(plane_w, plane_h) * 0.22) * repeat_scale
    repeat_h = repeat_w * (th / tw)

    ys, xs = np.nonzero(mask)
    xs_f, ys_f = xs.astype(np.float64), ys.astype(np.float64)

    def to_plane(px, py):
        z = H[2, 0] * px + H[2, 1] * py + H[2, 2]
        z = np.where(np.abs(z) < 1e-6, 1e-6, z)
        return (
            (H[0, 0] * px + H[0, 1] * py + H[0, 2]) / z,
            (H[1, 0] * px + H[1, 1] * py + H[1, 2]) / z,
        )

    fx, fy = to_plane(xs_f, ys_f)
    fx1, fy1 = to_plane(xs_f + 1, ys_f)
    fx2, fy2 = to_plane(xs_f, ys_f + 1)

    def rot(ax, ay):
        dx = ax - plane_cx
        dy = ay - plane_cy
        return dx * cos - dy * sin, dx * sin + dy * cos

    rx, ry = rot(fx, fy)
    rx1, ry1 = rot(fx1, fy1)
    rx2, ry2 = rot(fx2, fy2)

    tcx, tcy = (rx / repeat_w) * tw, (ry / repeat_h) * th
    du = np.hypot((rx1 / repeat_w) * tw - tcx, (ry1 / repeat_h) * th - tcy)
    dv = np.hypot((rx2 / repeat_w) * tw - tcx, (ry2 / repeat_h) * th - tcy)
    footprint = np.maximum(np.maximum(du, dv), 1e-3)
    lod = np.log2(footprint) + 0.5

    if lay_mode == "cells":
        u, v, cell_tone, grout_blend = procedural_cells(
            rx, ry, repeat_w, repeat_h, grout_frac, tw, th, footprint
        )
    else:
        u = np.mod(rx / repeat_w, 1.0)
        v = np.mod(ry / repeat_h, 1.0)
    l0 = np.clip(np.floor(lod), 0, max_l).astype(np.int64)
    frac = np.clip(lod - l0, 0, 1)

    sample = np.zeros((len(xs), 3), np.float64)
    for lev in range(max_l + 1):
        sel = l0 == lev
        if not sel.any():
            continue
        a = mips[lev]
        sa = sample_bilinear_wrap(a, u[sel] * a.shape[1], v[sel] * a.shape[0])
        if lev < max_l:
            b = mips[lev + 1]
            sb = sample_bilinear_wrap(b, u[sel] * b.shape[1], v[sel] * b.shape[0])
            sample[sel] = sa + (sb - sa) * frac[sel][:, None]
        else:
            sample[sel] = sa

    if lay_mode == "cells":
        sample = (
            sample * cell_tone[:, None] * (1.0 - grout_blend[:, None])
            + grout_col[None, :] * grout_blend[:, None]
        )

    shade = (
        shade_lo + (shade_map[ys, xs] / 255.0) * (shade_hi - shade_lo)
        if shade_map is not None
        else np.full(len(xs), 1.0)
    )

    specular = np.zeros(len(xs))
    if lv:
        lvx, lvy = lv.get("x", 0.0), lv.get("y", 0.0)
        a = to_plane(np.array([w * 0.5]), np.array([h * 0.75]))
        step = min(w, h) * 0.05
        b = to_plane(np.array([w * 0.5 + lvx * step]), np.array([h * 0.75 + lvy * step]))
        dxv, dyv = b[0][0] - a[0][0], b[1][0] - a[1][0]
        ln = np.hypot(dxv, dyv)
        if ln > 1e-6:
            lvx, lvy = dxv / ln, dyv / ln
            dfx = (fx - plane_cx) / (plane_w * 0.5)
            dfy = (fy - plane_cy) / (plane_h * 0.5)
            dlen = np.hypot(dfx, dfy)
            ok = dlen > 0.01
            dot = np.where(ok, (dfx * lvx + dfy * lvy) / np.maximum(dlen, 1e-9), 0.0)
            specular = 0.12 * gloss * np.maximum(0, dot) ** 4

    texel = soft_clip(apply_shade(sample, shade[:, None] * col[None, :]) + specular[:, None] * 255.0)
    if finish == "gloss":
        texel = apply_reflection(texel, base, xs, ys, h, w)
    alpha = conf[ys, xs][:, None] if conf is not None else np.ones((len(xs), 1))

    out = base.copy()
    out[ys, xs] = np.clip(texel * alpha + base[ys, xs] * (1 - alpha), 0, 255)
    img = Image.fromarray(out.astype(np.uint8))
    scale = min(OUT_MAX_DIM / max(img.size), 1.0)
    if scale < 1.0:
        img = img.resize((round(img.width * scale), round(img.height * scale)), Image.BILINEAR)
    return img


def main():
    if len(sys.argv) != 4:
        print(__doc__)
        return 2
    render(sys.argv[1], sys.argv[2]).save(sys.argv[3])
    print(f"saved {sys.argv[3]}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())