"""N1 — horizontal band seams on periodic (geometric) tile assets.

Reproduces the canvas-engine texture path on a steep look-down floor
(rooms 4/5 of the June-10 test set) for three texture preparations:

  raw       : tile the source as-is (what "wrap" mode does)
  committed : makeSeamless v2 masked-shift (current organic path)
  snapped   : proposed period-snap crop — detect the pattern's x/y period by
              autocorrelation, crop to an integer number of periods, and let
              it wrap exactly. Falls back to masked-shift when no strong
              period exists (wood, stone), so organic sources are untouched.

Pass criteria:
  1. checkered.jpeg classifies "organic" (it is — 2.5 x 1.65 periods).
  2. Period-snap finds a period; the cropped tile re-classifies as "wrap".
  3. Seam-spike score (max row-to-row jump / median) of the snapped render
     is at or below the committed render's, and shows no spike rows.
  4. rustic-wood.jpg (organic) finds NO period -> falls back unchanged.
"""

import os
import numpy as np
from PIL import Image

HERE = os.path.dirname(os.path.abspath(__file__))
TILES = os.path.join(HERE, "..", "..", "frontend", "viz2d-demo", "src", "assets", "tiles")
OUT = os.path.join(HERE, "verify_out")
os.makedirs(OUT, exist_ok=True)


# ---------------------------------------------------------------- engine ports
def detect_wrap_mode(img):
    """Exact port of detectWrapMode (canvas-engine.ts)."""
    h, w, _ = img.shape
    if w < 4 or h < 4:
        return "organic", np.inf, np.inf

    def col_diff(xa, xb):
        return float(np.mean(np.abs(img[:, xa].astype(np.float64) - img[:, xb].astype(np.float64))))

    def row_diff(ya, yb):
        return float(np.mean(np.abs(img[ya].astype(np.float64) - img[yb].astype(np.float64))))

    mid_x, mid_y = w >> 1, h >> 1
    internal = max((col_diff(mid_x, mid_x + 1) + row_diff(mid_y, mid_y + 1)) / 2, 1.5)
    seam_x = col_diff(w - 1, 0) / internal
    seam_y = row_diff(h - 1, 0) / internal
    mode = "wrap" if (seam_x < 3 and seam_y < 3) else "organic"
    return mode, seam_x, seam_y


def make_seamless(img):
    """Exact port of makeSeamless v2 (masked shift)."""
    h, w, _ = img.shape
    half_w, half_h = w >> 1, h >> 1
    band_x = max(2, round(w * 0.12))
    band_y = max(2, round(h * 0.12))

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

    m = np.outer(edge(h, band_y), edge(w, band_x))[..., None]
    shifted = np.roll(img, (-half_h, -half_w), axis=(0, 1)).astype(np.float64)
    return (img.astype(np.float64) * m + shifted * (1 - m)).astype(np.uint8)


def find_period(img, axis):
    """Proposed: dominant pattern period along an axis (x: axis=1, y: axis=0).

    Mean-abs-diff between the image and itself shifted by each candidate lag;
    a strongly periodic texture has a deep minimum at the period. Searches a
    <=320px downsample, then refines at full resolution. Returns the full-res
    period or None.
    """
    raw = np.mean(img.astype(np.float64), axis=2)
    h, w = raw.shape
    # High-pass: subtract a 16px low-pass reconstruction so the photo's
    # lighting gradient doesn't put a floor under the diff at true alignment.
    lp = np.asarray(Image.fromarray(raw).resize((16, 16), Image.BILINEAR).resize((w, h), Image.BILINEAR))
    gray = raw - lp
    full = gray.shape[1] if axis == 1 else gray.shape[0]
    scale = max(1, int(np.ceil(full / 320)))
    ds = gray[::scale, ::scale]
    n = ds.shape[1] if axis == 1 else ds.shape[0]

    lags = np.arange(max(12, int(n * 0.18)), int(n * 0.80))
    if len(lags) < 4:
        return None
    diffs = []
    for lag in lags:
        if axis == 1:
            d = np.mean(np.abs(ds[:, lag:] - ds[:, :-lag]))
        else:
            d = np.mean(np.abs(ds[lag:] - ds[:-lag]))
        diffs.append(d)
    diffs = np.asarray(diffs)

    med = np.median(diffs) + 1e-6
    best = diffs.min()
    # Generous gate: the post-crop wrap re-classification is the real safety
    # net; this only needs to exclude clearly aperiodic textures (wood ~0.85+,
    # stone ~0.96 vs checker ~0.46).
    if best / med > 0.65:
        return None
    # fundamental period: smallest lag within 15% of the global minimum
    idx = np.nonzero(diffs <= best * 1.15)[0][0]
    coarse = int(lags[idx]) * scale

    # refine at full resolution
    lo = max(8, coarse - scale - 2)
    hi = min(full - 1, coarse + scale + 2)
    best_lag, best_d = None, np.inf
    for lag in range(lo, hi + 1):
        if axis == 1:
            d = np.mean(np.abs(gray[:, lag:] - gray[:, :-lag]))
        else:
            d = np.mean(np.abs(gray[lag:] - gray[:-lag]))
        if d < best_d:
            best_d, best_lag = d, lag
    return best_lag


def best_crop(img, k_periods, period, axis):
    """Micro-search +/-5px around k*period for the crop length L that tiles
    best. The crop [0, L) wraps perfectly iff line L (the true continuation
    in the source) matches line 0 — so that is the diff to minimize."""
    n = img.shape[1] if axis == 1 else img.shape[0]
    target = k_periods * period
    win = max(4, min(16, period // 4))  # window beats flat-cell degeneracy
    cands = []
    for cand in range(max(8, target - 5), min(n - win, target + 5) + 1):
        if axis == 1:
            d = np.mean(np.abs(img[:, cand:cand + win].astype(np.float64)
                               - img[:, 0:win].astype(np.float64)))
        else:
            d = np.mean(np.abs(img[cand:cand + win].astype(np.float64)
                               - img[0:win].astype(np.float64)))
        cands.append((cand, d))
    if not cands:
        return target
    dmin = min(d for _, d in cands)
    # among near-ties, prefer the length closest to exactly k periods
    near = [c for c, d in cands if d <= dmin * 1.15 + 1e-6]
    return min(near, key=lambda c: abs(c - target))


def seam_vs_period_pair(img, crop_len, k, period, axis):
    """Wrap-seam diff relative to the texture's own k-period-pair diff.

    The wrap seam joins content k periods apart, so the fair yardstick is two
    interior lines k periods apart in the ORIGINAL image: they match in
    structure (grout geometry) but differ by per-tile surface noise and
    accumulated photo-perspective drift — the achievable floor for this seam.
    """
    a = img.astype(np.float64)
    if axis == 0:
        a = a.transpose(1, 0, 2)  # treat rows as columns; jog becomes vertical
    n = a.shape[1]
    dist = min(k * period, n - 1)
    a0 = max(0, (n - 1 - dist) // 2)  # internal pair, centered, same distance
    win = max(2, min(8, period // 8))
    jog = max(2, round(period * 0.04))  # photo-perspective drift tolerance

    def pair_diff(xa, xb):
        # min over a small perpendicular jog: a slightly skewed lattice still
        # tiles acceptably (the jog reads as installation tolerance, not a seam)
        best = np.inf
        pa = a[:, xa:xa + win]
        for dy in range(-jog, jog + 1):
            if dy >= 0:
                d = np.mean(np.abs(pa[dy:] - a[: a.shape[0] - dy, xb:xb + win]))
            else:
                d = np.mean(np.abs(pa[:dy] - a[-dy:, xb:xb + win]))
            best = min(best, d)
        return best

    seam = pair_diff(crop_len, 0)
    # floor = median over several internal pairs at the same distance — a
    # single pair can land in unrepresentatively flat or busy content
    starts = range(0, n - dist - win, max(1, (n - dist - win) // 8 or 1))
    floors = [pair_diff(s + dist, s) for s in starts] or [pair_diff(a0 + dist, a0)]
    floor_ = float(np.median(floors))
    return seam / max(floor_, 1e-6)


def period_snap(img):
    """Proposed prepare step: crop to integer periods if strongly periodic and
    the seam is no worse than the texture's own period-pair noise floor;
    otherwise fall back to masked shift."""
    h, w, _ = img.shape
    px = find_period(img, axis=1)
    py = find_period(img, axis=0)
    if px and py:
        kx, ky = w // px, h // py
        # the continuation line (k*period) must exist in the source to verify
        # the wrap; an exact-multiple image keeps one period fewer
        while kx > 1 and kx * px >= w:
            kx -= 1
        while ky > 1 and ky * py >= h:
            ky -= 1
        if kx * px < w and ky * py < h \
                and kx >= 1 and ky >= 1 and kx * px >= 0.4 * w and ky * py >= 0.4 * h:
            cw = best_crop(img, kx, px, axis=1)
            ch = best_crop(img, ky, py, axis=0)
            rx = seam_vs_period_pair(img, cw, kx, px, axis=1)
            ry = seam_vs_period_pair(img, ch, ky, py, axis=0)
            if rx < 1.5 and ry < 1.5:
                return img[:ch, :cw], ("snap", px, py, rx, ry)
    return make_seamless(img), ("fallback", px, py, None, None)


def flatten_luminance(img):
    """R2-3b — remove the source photo's baked lighting so the texture behaves
    like albedo: divide by a heavily-blurred (toroidal) luminance field
    normalized to the texture mean. The blur wraps, so the correction is
    itself seamless; only the lighting gradient is equalized, grain/veining
    survives. Mirror of flattenLuminance (canvas-engine.ts) — keep in
    lockstep."""
    h, w, _ = img.shape
    lum = (
        img[..., 0] * 0.299 + img[..., 1] * 0.587 + img[..., 2] * 0.114
    ).astype(np.float64)

    def box_wrap(a, r, axis):
        win = 2 * r + 1
        out = np.zeros_like(a)
        for d in range(-r, r + 1):
            out += np.roll(a, -d, axis=axis)
        return out / win

    rx = max(2, w // 4)
    ry = max(2, h // 4)
    for _ in range(2):
        lum = box_wrap(lum, rx, axis=1)
        lum = box_wrap(lum, ry, axis=0)
    gain = np.clip(lum.mean() / np.maximum(lum, 1.0), 0.6, 1.6)
    out = np.clip(np.round(img.astype(np.float64) * gain[..., None]), 0, 255)
    return out.astype(np.uint8)


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


def sample_bilinear_wrap(level, x, y):
    h, w, _ = level.shape
    x0 = np.clip(np.floor(x), 0, w - 1).astype(np.int64)
    y0 = np.clip(np.floor(y), 0, h - 1).astype(np.int64)
    x1 = (x0 + 1) % w
    y1 = (y0 + 1) % h
    fx = (x - np.floor(x))[..., None]
    fy = (y - np.floor(y))[..., None]
    p00, p10 = level[y0, x0], level[y0, x1]
    p01, p11 = level[y1, x0], level[y1, x1]
    return (p00 * (1 - fx) * (1 - fy) + p10 * fx * (1 - fy)
            + p01 * (1 - fx) * fy + p11 * fx * fy)


def render_floor(tex, img_w=900, img_h=700, repeat_w=180.0):
    """Steep look-down floor like rooms 4/5: asymmetric trapezoid -> deep plane."""
    th, tw, _ = tex.shape
    repeat_h = repeat_w * (th / tw)
    plane_w, plane_h = 900.0, 1600.0

    # image trapezoid (slight asymmetry = synthetic-VP shear) -> plane rect
    src = np.array([[260, 120], [610, 120], [900, 700], [0, 700]], np.float64)
    dst = np.array([[0, 0], [plane_w, 0], [plane_w, plane_h], [0, plane_h]], np.float64)
    A = []
    for (sx, sy), (dx_, dy_) in zip(src, dst):
        A.append([sx, sy, 1, 0, 0, 0, -dx_ * sx, -dx_ * sy])
        A.append([0, 0, 0, sx, sy, 1, -dy_ * sx, -dy_ * sy])
    b = dst.reshape(-1)
    hvec = np.linalg.solve(np.asarray(A), b)
    H = np.append(hvec, 1).reshape(3, 3)

    xs, ys = np.meshgrid(np.arange(img_w, dtype=np.float64), np.arange(img_h, dtype=np.float64))

    def to_plane(px, py):
        zz = H[2, 0] * px + H[2, 1] * py + H[2, 2]
        return ((H[0, 0] * px + H[0, 1] * py + H[0, 2]) / zz,
                (H[1, 0] * px + H[1, 1] * py + H[1, 2]) / zz)

    fx, fy = to_plane(xs, ys)
    fx1, fy1 = to_plane(xs + 1, ys)
    fx2, fy2 = to_plane(xs, ys + 1)

    # floor mask: inside the plane rect
    mask = (fx >= 0) & (fx < plane_w) & (fy >= 0) & (fy < plane_h)

    u = np.mod(fx / repeat_w, 1.0)
    v = np.mod(fy / repeat_h, 1.0)
    tcx, tcy = (fx / repeat_w) * tw, (fy / repeat_h) * th
    du = np.hypot((fx1 / repeat_w) * tw - tcx, (fy1 / repeat_h) * th - tcy)
    dv = np.hypot((fx2 / repeat_w) * tw - tcx, (fy2 / repeat_h) * th - 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), 0, max_l).astype(np.int64)
    f = np.clip(lod - l0, 0, 1)

    out = np.zeros((img_h, img_w, 3), np.float64)
    for lev in range(max_l + 1):
        sel = mask & (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])
        fb = f[sel][..., None]
        if lev < max_l:
            bl = mips[lev + 1]
            sb = sample_bilinear_wrap(bl, u[sel] * bl.shape[1], v[sel] * bl.shape[0])
            out[sel] = sa + (sb - sa) * fb
        else:
            out[sel] = sa
    return out.astype(np.uint8), mask


def sharpness_profile(render, mask):
    """Per-row mean |horizontal gradient| inside the floor. Crossfade ghost
    bands collapse local contrast, so they show up as dips in this profile.
    Comparing per-row against the raw render cancels the natural LOD falloff."""
    g = np.mean(render.astype(np.float64), axis=2)
    grad = np.abs(np.diff(g, axis=1))
    both = mask[:, :-1] & mask[:, 1:]
    rows = np.nonzero(both.sum(axis=1) > 200)[0]
    prof = np.array([np.mean(grad[y, both[y]]) for y in rows])
    return rows, prof


# -------------------------------------------------------------------- run
def load(name):
    return np.asarray(Image.open(os.path.join(TILES, name)).convert("RGB"))


def main():
    ok = True

    checker = load("checkered.jpeg")
    mode, sx, sy = detect_wrap_mode(checker)
    print(f"checkered.jpeg  : {checker.shape[1]}x{checker.shape[0]}  mode={mode}  "
          f"seamX={sx:.1f}x  seamY={sy:.1f}x  (threshold 3x)")
    if mode != "organic":
        print("  !! expected organic"); ok = False

    snapped, info = period_snap(checker)
    tag, px, py, csx, csy = info
    print(f"period-snap     : {tag}  periodX={px} periodY={py}  "
          f"crop={snapped.shape[1]}x{snapped.shape[0]}  "
          f"crop-seamX={csx if csx is None else f'{csx:.1f}x'} "
          f"crop-seamY={csy if csy is None else f'{csy:.1f}x'}")
    if tag != "snap":
        print("  !! period-snap did not engage on the checker"); ok = False

    healed = make_seamless(checker)

    # Real-photo renders: saved for visual / golden-image comparison (the
    # photo's per-tile texture variance defeats simple numeric seam metrics).
    base_repeat = 180.0
    snap_repeat = base_repeat * (snapped.shape[1] / checker.shape[1])
    for name, tex, rep in [("raw", checker, base_repeat),
                           ("committed", healed, base_repeat),
                           ("snapped", snapped, snap_repeat)]:
        render, _ = render_floor(tex, repeat_w=rep)
        Image.fromarray(render).save(os.path.join(OUT, f"n1_{name}.png"))
    print(f"photo renders saved to verify_out/n1_*.png (visual check)")

    # ---- synthetic certification: ground truth is known exactly ----------
    # Perfect checker, period 64, sized to a NON-integer period count so it
    # classifies organic. The ideal result is the exact 5-period crop tiled
    # raw; the snapped pipeline must reproduce it pixel-for-pixel.
    per = 64
    sw, sh = per * 5 + 37, per * 4 + 21
    yy, xx = np.mgrid[0:sh, 0:sw]
    cells = ((xx // (per // 2)) + (yy // (per // 2))) % 2
    rng = np.random.default_rng(7)
    noise = rng.normal(0, 4, (sh, sw))
    synth = np.stack([np.where(cells, 205, 120) + noise,
                      np.where(cells, 200, 90) + noise,
                      np.where(cells, 190, 60) + noise], axis=2).clip(0, 255).astype(np.uint8)

    smode, ssx, ssy = detect_wrap_mode(synth)
    print(f"synthetic checker: {sw}x{sh} period={per}  mode={smode}  "
          f"seam=({ssx:.1f}x,{ssy:.1f}x)")
    if smode != "organic":
        print("  !! synthetic checker should classify organic"); ok = False

    s_snap, s_info = period_snap(synth)
    print(f"synthetic snap   : {s_info[0]}  periodX={s_info[1]} periodY={s_info[2]}  "
          f"crop={s_snap.shape[1]}x{s_snap.shape[0]}")
    if s_info[0] != "snap":
        print("  !! period-snap did not engage on synthetic checker"); ok = False
    else:
        if s_info[1] % per != 0 or s_info[2] % per != 0:
            print(f"  !! found period not a multiple of {per}"); ok = False

        ideal = synth[:(sh // per) * per, :(sw // per) * per]
        rep_snap = 180.0 * (s_snap.shape[1] / sw)
        rep_ideal = 180.0 * (ideal.shape[1] / sw)
        r_snap, m1 = render_floor(s_snap, repeat_w=rep_snap)
        r_ideal, m2 = render_floor(ideal, repeat_w=rep_ideal)
        r_healed, _ = render_floor(make_seamless(synth), repeat_w=180.0)
        Image.fromarray(r_snap).save(os.path.join(OUT, "n1_synth_snapped.png"))
        Image.fromarray(r_ideal).save(os.path.join(OUT, "n1_synth_ideal.png"))
        Image.fromarray(r_healed).save(os.path.join(OUT, "n1_synth_committed.png"))

        m = m1 & m2
        d_snap = float(np.mean(np.abs(r_snap[m].astype(float) - r_ideal[m].astype(float))))
        d_healed = float(np.mean(np.abs(r_healed[m].astype(float) - r_ideal[m].astype(float))))
        print(f"synthetic render : snapped-vs-ideal={d_snap:.2f}  "
              f"committed-vs-ideal={d_healed:.2f}  (mean abs px)")
        if d_snap > 3.0:
            print("  !! snapped render deviates from ideal"); ok = False
        if d_healed < d_snap:
            print("  !! masked-shift unexpectedly beats period-snap"); ok = False

    # no-regression: organic wood must NOT engage period-snap
    wood = load("rustic-wood.jpg")
    _, winfo = period_snap(wood)
    print(f"rustic-wood.jpg : period-snap -> {winfo[0]}  "
          f"(periodX={winfo[1]} periodY={winfo[2]})")
    if winfo[0] != "fallback":
        print("  !! wood should fall back to masked-shift"); ok = False

    # informational: how do the other catalog tiles classify?
    for name in ["floor-natural-stone.jpg", "basalt-outside-wal.jpg", "mosaic-tile.jpg"]:
        t = load(name)
        m, a, b = detect_wrap_mode(t)
        _, i2 = period_snap(t)
        print(f"{name:24s}: mode={m:7s} seam=({a:.1f}x,{b:.1f}x) snap={i2[0]} px={i2[1]} py={i2[2]}")

    print("\n" + ("ALL N1 CHECKS PASSED" if ok else "N1 CHECKS FAILED"))
    return 0 if ok else 1


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