#!/usr/bin/env python3
"""GUTS texture post: tileable -> grayscale -> Sobel normals -> WebP (Lane D).

Pure numpy + PIL. NO model, NO network, NO spend — this is arithmetic (PIPELINE.md §Textures.3).

Why post-processing exists: FLUX is *told* "seamless tiling texture" and does not obey — its
output has hard edge seams. Two exact fixes are implemented; `--tile-mode` picks:

  blend  (default)  4-way cosine partition-of-unity of the image's four half-rolls.
                    w00=sx*sy, w10=cx*sy, w01=sx*cy, w11=cx*cy with sx=sin²(πx/N), cx=1-sx.
                    Weights sum to 1 everywhere and are N-periodic, so the result tiles
                    EXACTLY; each copy's seam sits where its own weight is 0, so no seam
                    survives. Cost: quadrant centres blend all four copies -> softer detail.
  mirror            2N mirror-quad downscaled to N. Crisper, but visibly symmetric
                    ("butterfly"). PIPELINE.md sanctions this as the fallback fix.
  none              passthrough (for already-tiling sources / debugging).

  <mflux-py> pipeline/derive_maps.py --contact          # contact sheets only, no harvest
  <mflux-py> pipeline/derive_maps.py --harvest          # .genraw/*.png -> web/assets/gen/*.webp
  <mflux-py> pipeline/derive_maps.py --harvest --tile-mode mirror --only esophagus
"""
import json, math, os, sys
import numpy as np
from PIL import Image, ImageFilter

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
RAW = os.path.join(ROOT, "pipeline", ".genraw")
GEN = os.path.join(ROOT, "web", "assets", "gen")
SHOTS = os.path.join(ROOT, "docs", "shots", "laneD")
BATCH = os.path.join(ROOT, "pipeline", "batch_textures.json")
os.makedirs(GEN, exist_ok=True)
os.makedirs(SHOTS, exist_ok=True)

WEBP_Q = 82          # walls: lossy is fine, they are luminance/AO detail
WEBP_Q_N = 90        # normals: lossy banding shows in the lighting, so spend a little more
NORMAL_STRENGTH = 2.4
OUT_SIZE = 1024      # house law: WebP <= 1024 (TECH.md)
MATCAP_SIZE = 512

# Pack coherence: FLUX gives every prompt its own exposure (the esophagus came back ~2 stops
# darker than the villi). These are DETAIL maps that get MULTIPLIED into a biome tint by Lane
# A's shader, so a dark texture doesn't read as "dark tissue", it reads as "this biome's tint
# is broken". Normalising every wall to the same mean/spread makes the tint mean the same
# thing in every biome; the shader, not the texture, decides brightness.
LEVELS_MEAN = 0.50
LEVELS_STD = 0.19


# ── image <-> array ─────────────────────────────────────────────────────────────────────
def load_gray(path):
    """Load as float32 luminance 0..1. Textures are authored GRAYSCALE — the biome tint is
    Lane A's shader's job (ART_BIBLE), so colour here would fight the tint."""
    im = Image.open(path).convert("L")
    return np.asarray(im, dtype=np.float32) / 255.0


def save_gray(a, path, q=WEBP_Q):
    im = Image.fromarray(np.clip(a * 255.0, 0, 255).astype(np.uint8), mode="L")
    im.save(path, "WEBP", quality=q, method=6)
    return os.path.getsize(path)


def save_rgb(a, path, q=WEBP_Q_N):
    im = Image.fromarray(np.clip(a * 255.0, 0, 255).astype(np.uint8), mode="RGB")
    im.save(path, "WEBP", quality=q, method=6)
    return os.path.getsize(path)


def resize(a, n):
    im = Image.fromarray(np.clip(a * 255.0, 0, 255).astype(np.uint8), mode="L")
    return np.asarray(im.resize((n, n), Image.LANCZOS), dtype=np.float32) / 255.0


# ── tiling ──────────────────────────────────────────────────────────────────────────────
def tile_blend(a):
    """4-way cosine partition-of-unity blend. Exactly tileable (see module docstring)."""
    h, w = a.shape
    x = np.arange(w, dtype=np.float32)
    y = np.arange(h, dtype=np.float32)
    sx = np.sin(np.pi * x / w) ** 2          # 0 at x=0 and x=w, 1 at x=w/2
    sy = np.sin(np.pi * y / h) ** 2
    cx, cy = 1.0 - sx, 1.0 - sy
    r_x = np.roll(a, w // 2, axis=1)
    r_y = np.roll(a, h // 2, axis=0)
    r_xy = np.roll(r_x, h // 2, axis=0)
    out = (a * np.outer(sy, sx) + r_x * np.outer(sy, cx)
           + r_y * np.outer(cy, sx) + r_xy * np.outer(cy, cx))
    # the 4-way average in the quadrant centres flattens contrast; pull it back to the
    # source's spread so the pack stays consistent under the shader's tint.
    return match_stats(out, a)


def tile_mirror(a):
    """Mirror quad -> downscale. Exactly tileable, but symmetric."""
    top = np.concatenate([a, a[:, ::-1]], axis=1)
    full = np.concatenate([top, top[::-1, :]], axis=0)
    return resize(full, a.shape[0])


def match_stats(x, ref):
    """Rescale x to ref's mean/std (contrast restore), then clip."""
    sx = x.std() or 1.0
    return np.clip((x - x.mean()) * (ref.std() / sx) + ref.mean(), 0.0, 1.0)


def normalize_levels(a, mean=LEVELS_MEAN, std=LEVELS_STD):
    """Bring a wall to the pack's common exposure. Linear (shape-preserving) — soft-clipped
    at the ends so the SEM speckle highlights roll off instead of blowing out to flat white."""
    sa = a.std() or 1.0
    x = (a - a.mean()) * (std / sa) + mean
    return np.clip(0.5 + np.tanh((x - 0.5) * 1.9) * 0.5, 0.0, 1.0)


def seam_error(a):
    """Mean |delta| across the wrap seams vs mean |delta| inside the image. ~1.0 = tiles
    perfectly (the seam is as smooth as ordinary interior detail); >>1 = visible seam.
    This is the NUMBER that decides tiling — not vibes."""
    inner_x = np.abs(np.diff(a, axis=1)).mean()
    inner_y = np.abs(np.diff(a, axis=0)).mean()
    seam_x = np.abs(a[:, 0] - a[:, -1]).mean()
    seam_y = np.abs(a[0, :] - a[-1, :]).mean()
    return (seam_x / (inner_x or 1e-6) + seam_y / (inner_y or 1e-6)) / 2


# ── normals ─────────────────────────────────────────────────────────────────────────────
def height_of(a):
    """Grayscale -> height via a small blur pyramid: keeps the big folds, drops the
    per-pixel SEM speckle that would otherwise become normal-map noise."""
    im = Image.fromarray((a * 255).astype(np.uint8), mode="L")
    b1 = np.asarray(im.filter(ImageFilter.GaussianBlur(1.0)), dtype=np.float32) / 255.0
    b2 = np.asarray(im.filter(ImageFilter.GaussianBlur(4.0)), dtype=np.float32) / 255.0
    return 0.65 * b1 + 0.35 * b2


def normal_of(a, strength=NORMAL_STRENGTH):
    """Sobel over a WRAPPED height field -> tangent-space normal map. np.roll wrapping is
    what keeps the normal map tileable alongside its albedo."""
    h = height_of(a)
    def sh(dy, dx):
        return np.roll(np.roll(h, dy, axis=0), dx, axis=1)
    gx = ((sh(-1, -1) + 2 * sh(0, -1) + sh(1, -1)) - (sh(-1, 1) + 2 * sh(0, 1) + sh(1, 1)))
    gy = ((sh(-1, -1) + 2 * sh(-1, 0) + sh(-1, 1)) - (sh(1, -1) + 2 * sh(1, 0) + sh(1, 1)))
    nx, ny, nz = -gx * strength, -gy * strength, np.ones_like(gx)
    ln = np.sqrt(nx * nx + ny * ny + nz * nz)
    return np.stack([(nx / ln) * 0.5 + 0.5, (ny / ln) * 0.5 + 0.5, (nz / ln) * 0.5 + 0.5], axis=-1)


# ── matcap ──────────────────────────────────────────────────────────────────────────────
def crop_matcap(a):
    """Find the lit ball on the black studio background and crop it square. A matcap must be
    the sphere touching all four edges or the lighting lookup is wrong."""
    m = a > (a.max() * 0.10)
    ys, xs = np.where(m)
    if len(xs) < 100:
        return resize(a, MATCAP_SIZE)
    x0, x1, y0, y1 = xs.min(), xs.max(), ys.min(), ys.max()
    cx, cy = (x0 + x1) / 2, (y0 + y1) / 2
    r = max(x1 - x0, y1 - y0) / 2
    h, w = a.shape
    x0 = int(max(0, cx - r)); x1 = int(min(w, cx + r))
    y0 = int(max(0, cy - r)); y1 = int(min(h, cy + r))
    return resize(a[y0:y1, x0:x1], MATCAP_SIZE)


# ── contact sheets (the eyeball law) ────────────────────────────────────────────────────
def contact_sheet(entries, path, cell=256):
    """One row per texture: [source] [tiled 2x2] [normal]. The 2x2 is the tile check — a seam
    shows up instantly as a cross through the middle."""
    if not entries:
        return None
    rows = len(entries)
    sheet = Image.new("RGB", (cell * 3 + 40, rows * cell + 20), (10, 12, 14))
    for i, (name, src, tiled, nrm) in enumerate(entries):
        y = 10 + i * cell
        def put(a, x, mode="L"):
            im = Image.fromarray(np.clip(a * 255, 0, 255).astype(np.uint8), mode=mode)
            sheet.paste(im.resize((cell, cell), Image.LANCZOS).convert("RGB"), (x, y))
        put(src, 10)
        two = np.concatenate([np.concatenate([tiled, tiled], axis=1)] * 2, axis=0)
        put(two, 20 + cell)
        put(nrm, 30 + cell * 2, mode="RGB")
    sheet.save(path, "PNG")
    return path


# ── harvest ─────────────────────────────────────────────────────────────────────────────
def harvest(only="", tile_mode="blend", contact_only=False):
    batch = json.load(open(BATCH)) if os.path.exists(BATCH) else {}
    srcs = sorted(f for f in os.listdir(RAW) if f.endswith(".png"))
    if only:
        srcs = [f for f in srcs if only in f]
    if not srcs:
        print(f"nothing in {RAW} matching *{only}* — run --local first"); return

    entries, results, wall_rows, cap_rows = [], {}, [], []
    for f in srcs:
        slug = os.path.splitext(f)[0]
        a = load_gray(os.path.join(RAW, f))
        is_cap = batch.get(slug, {}).get("kind") == "matcap" or slug.startswith("matcap")

        if is_cap:
            out = crop_matcap(a)
            kb_a = save_gray(out, os.path.join(GEN, f"{slug}.webp")) // 1024 if not contact_only else 0
            nrm = np.stack([out] * 3, axis=-1)      # sheet placeholder: matcaps get no normal
            cap_rows.append((slug, a, out, nrm))
            results[slug] = dict(kind="matcap", seam=None, kb=kb_a, size=MATCAP_SIZE)
            continue

        before = seam_error(a)
        mean_before = float(a.mean())
        tiled = {"blend": tile_blend, "mirror": tile_mirror, "none": lambda z: z}[tile_mode](a)
        tiled = normalize_levels(tiled)
        if tiled.shape[0] != OUT_SIZE:
            tiled = resize(tiled, OUT_SIZE)
        after = seam_error(tiled)
        nrm = normal_of(tiled)
        kb_a = kb_n = 0
        if not contact_only:
            kb_a = save_gray(tiled, os.path.join(GEN, f"{slug}.webp")) // 1024
            kb_n = save_rgb(nrm, os.path.join(GEN, f"{slug}_n.webp")) // 1024
        wall_rows.append((slug, a, tiled, nrm))
        results[slug] = dict(kind="wall", seam_before=round(float(before), 2),
                             seam_after=round(float(after), 2), kb=kb_a, kb_n=kb_n,
                             mean_before=round(mean_before, 3), mean_after=round(float(tiled.mean()), 3),
                             size=OUT_SIZE, tile_mode=tile_mode)
        print(f"  {slug:22s} seam {before:5.2f} -> {after:4.2f}   "
              f"mean {mean_before:.2f} -> {tiled.mean():.2f}   "
              f"{kb_a:>4}KB +{kb_n:>4}KB normal   [{tile_mode}]")

    for slug, r in results.items():
        if slug in batch:
            batch[slug].update({k: v for k, v in r.items() if k != "kind"})
    if batch:
        json.dump(batch, open(BATCH, "w"), indent=2, sort_keys=True)

    if wall_rows:
        p = contact_sheet(wall_rows, os.path.join(SHOTS, "contact_walls.png"))
        print(f"  contact sheet -> {p}")
    if cap_rows:
        p = contact_sheet(cap_rows, os.path.join(SHOTS, "contact_matcaps.png"))
        print(f"  contact sheet -> {p}")
    if not contact_only:
        tot = sum(r.get("kb", 0) + r.get("kb_n", 0) for r in results.values())
        print(f"harvested {len(results)} -> {GEN}  ({tot/1024:.2f}MB)  "
              f"(now run build_manifest.py)")
    return results


if __name__ == "__main__":
    only = sys.argv[sys.argv.index("--only") + 1] if "--only" in sys.argv else ""
    mode = sys.argv[sys.argv.index("--tile-mode") + 1] if "--tile-mode" in sys.argv else "blend"
    harvest(only=only, tile_mode=mode, contact_only="--contact" in sys.argv)