#!/usr/bin/env python3
"""
skirt_garment_weights.py — weight a skirt garment to the skirt_* spring ring.

    python clothing/skirt_garment_weights.py [garment.gltf]
    (default: assets/quaternius/outfits/kapahaka/Female_KapahakaSB_Legs.gltf)

Rewrites JOINTS_0/WEIGHTS_0 in place on a garment whose skin ALREADY carries the
SAME skirt ring as the body it is weighted against — both the strand count and the
segment count are read off that body, and a mismatch is a hard failure, not a
silent partial rebind (see skirt_rig_body.py). Geometry, materials,
UVs and the joint list are untouched — only the skin weights are recomputed, and
they are derived from vertex POSITION, so re-running is idempotent.

Why this exists: the first skirt-boned piupiu was weighted by an ad-hoc script
that wasn't kept, and it was badly lopsided — strand 1 got ZERO weight, strand 2
got 8x strand 6, and only 31% of the garment sat on the ring at all (measured
2026-07-31). The spring sim then moved a few strands and left the rest skinned to
the thigh, so the skirt hitched to one side instead of hanging. Weights are the
product here, not the sim tuning.

The three influences, per vertex:

  WAISTBAND (top few cm)  → pelvis alone. A waistband that swings looks broken;
                            it must ride the hip exactly as the body does.
  PANEL                   → the two nearest strands by AZIMUTH, blended linearly
                            across the 45 degree gap. Bone names carry no
                            directional meaning; ring positions do, so azimuths
                            are read off the bones rather than assumed from the
                            NN index. Uniform coverage is guaranteed by
                            construction — no strand can come out dead.
  THIGH FOLLOW            → mixed in by PROXIMITY to the thigh segment, so cloth
                            resting on the thigh is carried up WITH the knee and
                            keeps covering it. This is the difference between a
                            skirt that drapes over a raised knee and one the
                            spring flicks aside to bare the thigh. It fades out
                            over the bottom of the panel: the hem must stay free
                            for the spring to swing it, or the skirt turns into
                            a pair of trousers painted on the legs.

Godot takes only 4 influences per vertex, so the assembled set is truncated to
the 4 heaviest and renormalised.
"""
import json
import math
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ANIM_ROOT = os.path.dirname(HERE)
# This repo AUTHORS clothing; ariki-game only RECEIVES delivered outfits and bodies, so
# every asset path below points into the game checkout (a sibling of this one). Override
# with ARIKI_GAME_ROOT if the checkouts are not side by side.
GAME = os.environ.get("ARIKI_GAME_ROOT") or os.path.join(
    os.path.dirname(ANIM_ROOT), "ariki-game")
DEFAULT = os.path.join(
    GAME, "assets/quaternius/outfits/kapahaka/Female_KapahakaSB_Legs.gltf")
SRC = sys.argv[1] if len(sys.argv) > 1 else DEFAULT
BODY = sys.argv[2] if len(sys.argv) > 2 else os.path.join(
    GAME, "assets/quaternius/derived-bodies/Ariki_Female_QuatSkin_SkirtRig.glb")

# Strand + segment counts are READ OFF THE RIG (see body_ring in main), never constants.

# Everything ABOVE the ring top is pelvis by necessity — no skirt bone reaches it.
# On the piupiu that is already ~870 verts (the garment runs to the true waist at
# y=1.113 while the ring starts at the hip, y=0.952, where a piupiu is actually
# tied), so this only adds a thin band BELOW the ring top and must stay small or
# it freezes the top third of the panel.
BAND = 0.02
# Thigh follow is the one influence that can LIFT cloth: collision only ever pushes, but a
# thigh-weighted vertex is carried wherever the thigh goes. At 0.55 the swinging leg hauled
# the hem up with it — the skirt rode up on the upswing while the standing side hung long
# (back view, walk 2026-07-31). Keep just enough to drape cloth over a raised knee, and let
# collision do the rest now that it actually runs.
# ENV-OVERRIDABLE (defaults unchanged — an unset env behaves exactly as before).
# Added because these were tuned against the SHORT continuous piupiu, and a longer
# strand skirt needs different numbers: bottomTest1 hangs to y 0.451, so HEM_FREE 0.55
# puts its whole KNEE region (y~0.517) inside the no-follow band and the knee comes
# straight through the cloth. Per-garment tuning without touching the shared defaults.
def _envf(key, default):
    import os
    try:
        return float(os.environ[key])
    except (KeyError, ValueError):
        return default

FOLLOW_MAX = _envf("SKIRT_FOLLOW_MAX", 0.20)
CONTACT_R = _envf("SKIRT_CONTACT_R", 0.13)   # at/inside this distance from the thigh axis = full contact
FALLOFF = _envf("SKIRT_FALLOFF", 0.10)       # follow decays to 0 over this much extra distance
HEM_FREE = _envf("SKIRT_HEM_FREE", 0.55)     # bottom fraction of panel with NO thigh follow
MAX_INFLUENCES = 4

COMP_FMT = {5120: "b", 5121: "B", 5122: "h", 5123: "H", 5125: "I", 5126: "f"}
NCOMP = {"SCALAR": 1, "VEC2": 2, "VEC3": 3, "VEC4": 4, "MAT4": 16}


def log(msg):
    print(f"[skirt_weights] {msg}", flush=True)


def smoothstep(x):
    x = max(0.0, min(1.0, x))
    return x * x * (3.0 - 2.0 * x)


def mat4_inverse_translation(m):
    """Translation of inverse(m), where m is a glTF 4x4 — COLUMN-major, so the
    flat index for row r / col c is c*4+r (reading it row-major silently yields
    the transpose, which lands every bone at ~the origin). Gauss-Jordan rather
    than a rigid-transform shortcut because IBMs may carry scale.
    Returns (x, y, z) — the bone head in the space POSITION lives in."""
    a = [[m[c * 4 + r] for c in range(4)] + [1.0 if r == i else 0.0 for i in range(4)]
         for r in range(4)]
    for col in range(4):
        piv = max(range(col, 4), key=lambda r: abs(a[r][col]))
        if abs(a[piv][col]) < 1e-12:
            raise ValueError("singular inverse-bind matrix")
        a[col], a[piv] = a[piv], a[col]
        d = a[col][col]
        a[col] = [v / d for v in a[col]]
        for r in range(4):
            if r != col and a[r][col] != 0.0:
                f = a[r][col]
                a[r] = [v - f * w for v, w in zip(a[r], a[col])]
    return (a[0][7], a[1][7], a[2][7])   # inverse's 4th column = translation


def point_seg_distance(p, a, b):
    """Distance from p to segment ab (thigh head→knee)."""
    ab = tuple(b[i] - a[i] for i in range(3))
    ap = tuple(p[i] - a[i] for i in range(3))
    ll = sum(v * v for v in ab)
    t = 0.0 if ll < 1e-12 else max(0.0, min(1.0, sum(ap[i] * ab[i] for i in range(3)) / ll))
    closest = tuple(a[i] + ab[i] * t for i in range(3))
    return math.sqrt(sum((p[i] - closest[i]) ** 2 for i in range(3)))


class Gltf:
    def __init__(self, path):
        self.path = path
        if path.lower().endswith(".glb"):
            # GLB: 12-byte header, then JSON chunk, then the BIN chunk.
            with open(path, "rb") as f:
                blob = f.read()
            _, _, _ = struct.unpack_from("<III", blob, 0)
            clen, _ = struct.unpack_from("<II", blob, 12)
            self.d = json.loads(blob[20:20 + clen])
            blen, _ = struct.unpack_from("<II", blob, 20 + clen)
            start = 20 + clen + 8
            self.buf = bytearray(blob[start:start + blen])
            self.bin_path = None          # GLB is read-only here (body donor)
            return
        with open(path) as f:
            self.d = json.load(f)
        self.bin_path = os.path.join(os.path.dirname(path), self.d["buffers"][0]["uri"])
        with open(self.bin_path, "rb") as f:
            self.buf = bytearray(f.read())

    def read(self, idx):
        a = self.d["accessors"][idx]
        bv = self.d["bufferViews"][a["bufferView"]]
        off = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
        fmt = COMP_FMT[a["componentType"]]
        n = NCOMP[a["type"]]
        size = struct.calcsize("<" + fmt * n)
        stride = bv.get("byteStride") or size
        return [struct.unpack_from("<" + fmt * n, self.buf, off + i * stride)
                for i in range(a["count"])]

    def overwrite(self, idx, rows):
        """Rewrite an accessor's bytes IN PLACE, keeping its declared layout.

        Appending fresh bufferViews instead would orphan the old ones and grow
        the .bin on every run, so the tool would stop being safely re-runnable.
        Both skin attributes here own a tightly-packed dedicated bufferView, so
        an in-place write is exact — verified against the layout, not assumed."""
        a = self.d["accessors"][idx]
        bv = self.d["bufferViews"][a["bufferView"]]
        fmt = COMP_FMT[a["componentType"]]
        n = NCOMP[a["type"]]
        size = struct.calcsize("<" + fmt * n)
        if len(rows) != a["count"]:
            raise SystemExit(f"[skirt_weights] FATAL: accessor {idx} holds "
                             f"{a['count']} rows, got {len(rows)}")
        if bv.get("byteStride") not in (None, size) or bv["byteLength"] != size * len(rows):
            raise SystemExit(f"[skirt_weights] FATAL: accessor {idx} shares or "
                             "interleaves its bufferView — in-place write unsafe")
        if a["componentType"] == 5121 and any(v > 255 for row in rows for v in row):
            raise SystemExit("[skirt_weights] FATAL: joint index over 255 will not fit "
                             "the garment's unsigned-byte JOINTS_0")
        off = bv.get("byteOffset", 0) + a.get("byteOffset", 0)
        for i, row in enumerate(rows):
            struct.pack_into("<" + fmt * n, self.buf, off + i * size, *row)

    def save(self):
        self.d["buffers"][0]["byteLength"] = len(self.buf)
        with open(self.bin_path, "wb") as f:
            f.write(self.buf)
        with open(self.path, "w") as f:
            json.dump(self.d, f, separators=(",", ":"))


def main():
    g = Gltf(SRC)
    skin = g.d["skins"][0]
    joints = skin["joints"]
    names = [g.d["nodes"][j].get("name") for j in joints]
    slot = {n: i for i, n in enumerate(names)}

    for required in ("pelvis", "thigh_l", "thigh_r", "calf_l", "calf_r", "skirt_00"):
        if required not in slot:
            raise SystemExit(f"[skirt_weights] FATAL: no '{required}' joint in {SRC} "
                             "— is this a *_SkirtRig-skinned garment?")

    # ── REBIND the skirt bones from the body ────────────────────────────────────
    # The garment's inverseBindMatrices describe the ring it was armatured against.
    # Re-proportioning the ring in make_skirt_rig_body.py moves those bones, and a
    # garment still bound to the OLD ring is skinned to a skeleton that no longer
    # exists — it deforms toward stale rest positions. Names all still resolve, so no
    # name gate catches it; only the positions differ. Copy the body's skirt-bone
    # binds over so the garment is bound to the ring it will actually be worn on.
    body = Gltf(BODY)
    body_skin = body.d["skins"][0]
    body_names = [body.d["nodes"][j].get("name") for j in body_skin["joints"]]
    body_ibms = body.read(body_skin["inverseBindMatrices"])
    body_ibm = {n: m for n, m in zip(body_names, body_ibms) if n}

    ibms = g.read(skin["inverseBindMatrices"])

    # ── The RIG defines the ring; the garment must carry ALL of it ──────────────
    # STRAND count comes from the body, exactly like segment count. Hardcoding
    # either lets a garment silently keep an old ring. Measured failure: a 32-chain
    # body against an 8-chain garment rebound only the 8 names that happened to
    # match, then weighted a ring the mesh was never fitted to — 46:1 strand
    # imbalance, garbage on the character, and NO gate fired, because the old gate
    # compared the garment's bones against the garment's own joint list (2026-07-31).
    body_ring = {}
    for nm in body_ibm:
        if not nm.startswith("skirt_"):
            continue
        parts = nm.split("_")
        if len(parts) < 2 or not parts[1].isdigit():
            continue
        body_ring.setdefault(int(parts[1]), set()).add(nm)
    if not body_ring:
        raise SystemExit("[skirt_weights] FATAL: no skirt_* bones in "
                         f"{os.path.basename(BODY)} — is it a *_SkirtRig body?")

    absent = sorted(b for bones in body_ring.values() for b in bones if b not in slot)
    if absent:
        raise SystemExit(
            f"[skirt_weights] FATAL: {os.path.basename(BODY)} has "
            f"{len(body_ring)} strands / {sum(len(v) for v in body_ring.values())} skirt "
            f"bones, but this garment's skin cannot reference {len(absent)} of them "
            f"(e.g. {absent[:4]}).\nA garment can only be weighted to a ring its SKIN "
            "carries. Re-export the garment against this body first (the tailor lane's "
            "pipeline picks up every armature bone automatically), then re-run this.")
    stale = sorted(n for n in names
                   if n and n.startswith("skirt_") and n not in body_ibm)
    if stale:
        raise SystemExit(
            f"[skirt_weights] FATAL: this garment's skin carries {len(stale)} skirt bones "
            f"that do not exist on {os.path.basename(BODY)} (e.g. {stale[:4]}) — it was "
            "fitted to a DIFFERENT ring. Re-export against this body, or point --body at "
            "the one it was built for.")

    rebound, moved = 0, 0.0
    for i, nm in enumerate(names):
        if nm is None or not nm.startswith("skirt_") or nm not in body_ibm:
            continue
        before = mat4_inverse_translation(ibms[i])
        after = mat4_inverse_translation(body_ibm[nm])
        moved = max(moved, math.dist(before, after))
        ibms[i] = body_ibm[nm]
        rebound += 1
    log(f"rebound {rebound} skirt joints from {os.path.basename(BODY)} "
        f"({len(body_ring)} strands) — largest bind move {moved * 100:.1f} cm")
    # NOTE: the IBM write is deliberately AFTER the gates above. It used to run first,
    # so a rejected garment still had its bind matrices rewritten on the way out.
    g.overwrite(skin["inverseBindMatrices"], ibms)
    origin = {}
    for i, nm in enumerate(names):
        if nm is None:
            continue
        origin[nm] = mat4_inverse_translation(ibms[i])

    # glTF is Y-up: vertical is Y, the ring azimuth lives in the XZ plane.
    pelvis = origin["pelvis"]
    strands = []
    for k in sorted(body_ring):
        # Segment count comes from the rig, not a constant here — make_skirt_rig_body.py
        # owns N, and reading it back keeps the two tools from silently disagreeing.
        segs = [f"skirt_{k:02d}"]
        while f"{segs[0]}_{len(segs):02d}" in origin:
            segs.append(f"{segs[0]}_{len(segs):02d}")
        top = origin[segs[0]]
        # Segment HEAD heights are the real boundaries. Segment lengths are NOT uniform
        # (the rig deliberately uses a short upper segment), so nothing here may assume
        # span/N. The final segment simply owns everything below its head — the hem is a
        # bone TAIL, which glTF does not store, and deriving it is unnecessary.
        strands.append({
            "k": k, "segs": segs, "top_y": top[1],
            "heads": [origin[s][1] for s in segs],
            "az": math.atan2(top[2] - pelvis[2], top[0] - pelvis[0]),
        })
    strands.sort(key=lambda s: s["az"])
    n_seg = len(strands[0]["segs"])
    log(f"{n_seg} segments per strand (read off the rig), "
        f"boundaries y " + ", ".join(f"{y:.3f}" for y in strands[0]["heads"]))
    y_top = max(s["top_y"] for s in strands)
    ring_drop = y_top - min(s["heads"][-1] for s in strands)
    ring_radius = max(math.dist((s_top[0], s_top[2]), (pelvis[0], pelvis[2]))
                      for s_top in (origin[s["segs"][0]] for s in strands))
    log(f"ring: {len(strands)} strands, top y {y_top:.3f}, radius {ring_radius:.3f}, "
        f"last boundary {y_top - ring_drop:.3f}")
    if ring_drop < 0.02 or ring_radius < 0.03:
        raise SystemExit("[skirt_weights] FATAL: the ring has collapsed to a point, so "
                         "the bone origins are wrong — check the inverseBindMatrices "
                         "read (glTF matrices are COLUMN-major) before trusting weights")
    for s in strands:
        missing = [b for b in s["segs"] if b not in slot]
        if missing:
            raise SystemExit(f"[skirt_weights] FATAL: the rig has bones this garment's "
                             f"skin cannot reference: {missing}. A garment must be "
                             "re-armatured onto the current SkirtRig skeleton before it "
                             "can be weighted to it.")

    legs = [("thigh_l", origin["thigh_l"], origin["calf_l"]),
            ("thigh_r", origin["thigh_r"], origin["calf_r"])]

    prim = g.d["meshes"][0]["primitives"][0]
    pos = g.read(prim["attributes"]["POSITION"])
    # The vertical ramp (segment pick, thigh-follow hem fade) is measured against the
    # CLOTH, not the ring: the last bone's tail is the hem and glTF stores no tails, and
    # measuring to the last bone HEAD instead made the ramp collapse into the top 11 cm
    # the moment the ring stopped being an even 50/50 split.
    span = y_top - min(p[1] for p in pos)
    log(f"{os.path.basename(SRC)}: {len(pos)} verts, "
        f"cloth hangs {span:.3f} below the ring top")

    out_j, out_w = [], []
    stats = {"waistband": 0, "follow_sum": 0.0}
    per_strand = {k: 0.0 for k in body_ring}
    thigh_total = 0.0

    for v in pos:
        w = {}
        if v[1] > y_top - BAND:
            w[slot["pelvis"]] = 1.0
            stats["waistband"] += 1
        else:
            # ── azimuth → the two neighbouring strands, linear across the gap ──
            az = math.atan2(v[2] - pelvis[2], v[0] - pelvis[0])
            lo = None
            for i, s in enumerate(strands):
                nxt = strands[(i + 1) % len(strands)]
                d = (nxt["az"] - s["az"]) % (2 * math.pi)
                off = (az - s["az"]) % (2 * math.pi)
                if off <= d:
                    lo, hi, frac = s, nxt, (off / d if d > 1e-9 else 0.0)
                    break
            if lo is None:                      # numerically outside every gap
                lo = hi = min(strands, key=lambda s: abs((az - s["az"] + math.pi)
                                                         % (2 * math.pi) - math.pi))
                frac = 0.0
            share = {id(lo): (lo, 1.0 - frac)}
            if id(hi) in share:
                share[id(hi)] = (hi, share[id(hi)][1] + frac)
            else:
                share[id(hi)] = (hi, frac)

            # ── height → which segment(s) of the strand ──
            # Blend across the neighbouring pair rather than snapping, or the seam
            # between segments creases when the spring bends the strand.
            t = max(0.0, min(1.0, (y_top - v[1]) / span)) if span > 1e-9 else 0.0

            # ── thigh follow by proximity, faded out across the free hem ──
            near, dist = None, 1e9
            for bone, a, b in legs:
                dd = point_seg_distance(v, a, b)
                if dd < dist:
                    near, dist = bone, dd
            contact = 1.0 - smoothstep((dist - CONTACT_R) / FALLOFF)
            hem_fade = 1.0 - smoothstep((t - (1.0 - HEM_FREE)) / HEM_FREE)
            follow = FOLLOW_MAX * contact * hem_fade
            stats["follow_sum"] += follow

            if follow > 0.0:
                w[slot[near]] = w.get(slot[near], 0.0) + follow
            for strand, sh in share.values():
                ring = (1.0 - follow) * sh
                if ring <= 0.0:
                    continue
                segs, heads = strand["segs"], strand["heads"]
                # Continuous position in SEGMENT-INDEX units, from the real boundaries —
                # segment lengths are non-uniform, so t * len(segs) would misplace it.
                # Segment i spans [heads[i+1], heads[i]]; below the last head everything
                # belongs to the last bone.
                f = float(len(segs) - 1)
                for i in range(len(heads) - 1):
                    if v[1] > heads[i + 1]:
                        h = heads[i] - heads[i + 1]
                        f = i + (0.0 if h < 1e-9
                                 else max(0.0, min(1.0, (heads[i] - v[1]) / h)))
                        break
                i0 = min(len(segs) - 1, int(f))
                i1 = min(len(segs) - 1, i0 + 1)
                frac_seg = f - i0
                for si, part in ((i0, 1.0 - frac_seg), (i1, frac_seg)):
                    if part <= 0.0:
                        continue
                    s = slot[segs[si]]
                    w[s] = w.get(s, 0.0) + ring * part
                per_strand[strand["k"]] += ring
            thigh_total += follow

        # Godot reads 4 influences — keep the heaviest and renormalise.
        top4 = sorted(w.items(), key=lambda kv: -kv[1])[:MAX_INFLUENCES]
        total = sum(x for _, x in top4) or 1.0
        top4 = [(j, x / total) for j, x in top4]
        while len(top4) < MAX_INFLUENCES:
            top4.append((0, 0.0))
        out_j.append(tuple(j for j, _ in top4))
        out_w.append(tuple(x for _, x in top4))

    # ── gates FIRST: never leave a broken garment on disk ──
    n = len(pos)
    log(f"waistband (pelvis only): {stats['waistband']} verts")
    log(f"mean thigh follow on panel verts: "
        f"{stats['follow_sum'] / max(1, n - stats['waistband']):.3f}")
    log("per-strand ring weight (must be non-zero everywhere):")
    for k in sorted(per_strand):
        peak = max(per_strand.values())
        bar = "#" * int(per_strand[k] / peak * 40) if peak else ""
        log(f"   strand {k}: {per_strand[k]:7.1f}  {bar}")
    dead = [k for k in sorted(per_strand) if per_strand[k] <= 0.0]
    if dead:
        raise SystemExit(f"[skirt_weights] FATAL: strands {dead} got no weight — "
                         "the ring is not covered, the spring sim will tear the skirt")
    # Per-strand totals track VERTEX DENSITY per azimuth, not correctness — a
    # wrapped piupiu really does carry more geometry at the front overlap. The
    # gate that matters is the one above: every strand must be driving something,
    # or the spring sim pulls a panel the mesh can't follow and the skirt tears.
    log(f"heaviest/lightest strand ratio {max(per_strand.values()) / max(1e-9, min(per_strand.values())):.2f} "
        "(density, not a defect — the dead-strand gate above is the real check)")

    g.overwrite(prim["attributes"]["JOINTS_0"], out_j)
    g.overwrite(prim["attributes"]["WEIGHTS_0"], out_w)
    g.save()
    log(f"WROTE {SRC} + {os.path.basename(g.bin_path)} ({len(g.buf)} bytes)")


main()