"""Publish one solo-dataset clip after a codec round-trip, in the E2E turn viewer.

Diagnostic companion to `render_solo_gt_turn.py`. That script renders the
preprocessed GT rotations with NO codec and NO generator; this one takes the
SAME clip through `encoder -> continuous latent -> decoder` — the exact decode
contract the solo generator is trained against — and renders the result.

Because the two turns differ ONLY by the codec round-trip, comparing them in
the viewer separates two causes of bad arms:

    GT already wrong          -> preprocessing / rest-pose retargeting
    GT fine, round-trip wrong -> arms/fingers codec reconstruction

`face.json` and `audio.wav` are copied byte-for-byte from the GT turn rather
than re-rendered, so the body is provably the only difference. That also keeps
this process free of the face package, whose `scripts` module collides with
GestureVRM's (the same reason `infer_turn.py` runs the body in a subprocess).

WRITER (2026-08-05, ceiling-gate A/B) — `vrma_writer.write_vrma` is NOT used.
Its built-in generic skeleton has no upperChest, so that node is orphaned and
shoulders/neck rewire onto chest; the published pose breaks (elbow 72°->149°,
see render_solo_gt_turn.py docstring). The GT turn sidesteps this by shipping
the baked vrma verbatim — which a round-trip cannot do, since its rotations are
new. So we take that same baked vrma as the TEMPLATE and overwrite only the
rotation floats, leaving hierarchy, humanoid map and node rests byte-identical:

    npz basis = inv(rest_node) o channel     (vrma_to_npz.convert_humanoid)
    => channel = rest_node o basis           (this writer)

`--self-verify` (default on) re-writes the GT basis through the same path and
asserts it reproduces the baked channels, so the writer is proven neutral
before any codec output is judged. Measured 9.7e-08 max component error, i.e.
float32 noise. Bones outside the 50-bone contract (eyes, toes) keep their GT
channels; upperChest is absent from GG_11's humanoid map entirely.

    python render_solo_codec_roundtrip_turn.py \
        --npz /data/mocap-convertor/datasets/solo_20260722_body_v1/processed/npz/test/solo_20260722_174736_seg000126.npz \
        --gt-turn GT_15_restref_seg126 \
        --id CODEC_15_roundtrip_seg126
"""
import argparse
import json
import os
import shutil
import struct
import sys
import time

import numpy as np
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
KEMIX = os.path.abspath(os.path.join(HERE, ".."))
WH = os.path.join(KEMIX, "motion-blender", "experiments", "wave-hands")
# GestureVRM supplies the codec + 6D/quaternion conversions; motion-blender/lib
# supplies the glb reader. Deliberately no face import (see module docstring).
sys.path.insert(0, os.path.join(KEMIX, "motion", "GestureVRM"))
sys.path.insert(0, os.path.join(KEMIX, "motion-blender", "lib"))

import scripts.infer_vrma as iv                                  # noqa: E402
from dataloaders.vrm_dataset import quaternion_to_rotation_6d    # noqa: E402
from vrma_to_npz import FINGER_BONES, _accessor, _qmul, _read_glb  # noqa: E402

FPS = 30
RUN_ID = "solo_body_v1__50b_vel__scratch__r02"
RUN_ROOT = os.path.join(WH, "outputs", "gen_train", "runs", RUN_ID)

BODY20 = [
    "hips", "spine", "chest", "upperChest", "neck", "head",
    "leftShoulder", "leftUpperArm", "leftLowerArm", "leftHand",
    "rightShoulder", "rightUpperArm", "rightLowerArm", "rightHand",
    "leftUpperLeg", "leftLowerLeg", "leftFoot",
    "rightUpperLeg", "rightLowerLeg", "rightFoot",
]
# 손가락 본 수는 vrma_to_npz가 단일 소스 — 24본/30본 하드코딩 금지(2026-08-04
# 엄지 확장 때 이 파일이 44본에 묶여 있어 r02에서 못 쓰게 됐던 재발 방지).
BONES = BODY20 + list(FINGER_BONES)
NB = len(BONES)
NF = len(FINGER_BONES)

# (latent input dim, bone slice) — identical to validate_solo_body_codec.py.
PARTS = {
    "spine": (36, slice(0, 6)),
    "arms": (48, slice(6, 14)),
    "legs": (39, slice(14, 20)),
    "fingers": (NF * 6, slice(20, NB)),
}
CODEC_STRIDE = 4


def angular_error_deg(a, b):
    a = a / np.clip(np.linalg.norm(a, axis=-1, keepdims=True), 1e-8, None)
    b = b / np.clip(np.linalg.norm(b, axis=-1, keepdims=True), 1e-8, None)
    dots = np.abs(np.sum(a * b, axis=-1)).clip(0.0, 1.0)
    return np.degrees(2.0 * np.arccos(dots))


def load_codecs(device):
    """Load the selected scratch codecs, refusing any checkpoint outside this run."""
    selection = json.load(open(os.path.join(RUN_ROOT, "codec_selection.json")))
    if selection["run_id"] != RUN_ID:
        raise RuntimeError(f"foreign codec selection: {selection['run_id']}")
    meanstd = os.path.join(RUN_ROOT, "mean_std")

    models, stats = {}, {}
    for part, (dim, _) in PARTS.items():
        ckpt = os.path.abspath(selection["checkpoints"][part])
        if not ckpt.startswith(os.path.abspath(RUN_ROOT) + os.sep):
            raise RuntimeError(f"checkpoint escaped the isolated run: {ckpt}")
        model = iv.create_vqvae(dim)
        model.load_state_dict(
            torch.load(ckpt, map_location="cpu", weights_only=False)["net"])
        models[part] = model.eval().to(device)
        stats[part] = (
            np.load(os.path.join(meanstd, f"vrm_{part}_mean.npy")).astype(np.float32),
            np.load(os.path.join(meanstd, f"vrm_{part}_std.npy")).astype(np.float32),
        )
    return models, stats, selection


def roundtrip(q, trans, models, stats, device):
    """(T,NB,4) GT basis quats -> continuous codec round-trip -> (T,NB,4) + (T,3).

    The codec downsamples by CODEC_STRIDE, so T is edge-padded up to a multiple
    of it and trimmed back. The ceiling gate instead truncates; padding is used
    here so every published frame is codec output (a truncated tail spliced with
    GT would read as a pop at the loop boundary).
    """
    T = len(q)
    pad = (-T) % CODEC_STRIDE
    if pad:
        q = np.concatenate([q, np.repeat(q[-1:], pad, axis=0)], axis=0)
        trans = np.concatenate([trans, np.repeat(trans[-1:], pad, axis=0)], axis=0)
    frames = len(q)

    r6 = np.asarray(quaternion_to_rotation_6d(q)).astype(np.float32)
    trans_v = np.zeros_like(trans)
    trans_v[1:] = trans[1:] - trans[:-1]

    features = {
        "spine": r6[:, 0:6].reshape(frames, -1),
        "arms": r6[:, 6:14].reshape(frames, -1),
        "legs": np.concatenate([r6[:, 14:20].reshape(frames, -1), trans_v], axis=-1),
        "fingers": r6[:, 20:NB].reshape(frames, -1),
    }

    rec = {}
    with torch.no_grad():
        for part in PARTS:
            mean, std = stats[part]
            x = torch.from_numpy((features[part] - mean) / std).unsqueeze(0).to(device)
            latent = models[part].map2latent(x)
            out = models[part].decoder(latent.permute(0, 2, 1))[0].cpu().numpy()
            rec[part] = out * std + mean

    rec_r6 = np.zeros((frames, NB, 6), dtype=np.float32)
    rec_r6[:, 0:6] = rec["spine"].reshape(frames, 6, 6)
    rec_r6[:, 6:14] = rec["arms"].reshape(frames, 8, 6)
    rec_r6[:, 14:20] = rec["legs"][:, :36].reshape(frames, 6, 6)
    rec_r6[:, 20:NB] = rec["fingers"].reshape(frames, NF, 6)
    q_rec = np.asarray(iv.rotation_6d_to_quaternion(rec_r6))

    # 이동은 코덱이 속도로 복원 → GT 첫 프레임에서 적분(생성기 경로와 동일).
    rec_v = rec["legs"][:, 36:39]
    trans_rec = np.empty_like(trans)
    trans_rec[0] = trans[0]
    trans_rec[1:] = trans[0] + np.cumsum(rec_v[1:], axis=0)

    return q_rec[:T], trans_rec[:T]


def write_channels(template_vrma, basis, trans, out_path, bones=BONES):
    """Copy `template_vrma`, replacing only its rotation floats with `basis`.

    Returns the max |component| deviation the caller can use to self-verify.
    Bones absent from the humanoid map (upperChest on GG_11) or outside `bones`
    (eyes, toes) keep the template's own channels.
    """
    gltf, blob = _read_glb(template_vrma)
    blob = bytearray(blob)
    nodes, anim = gltf["nodes"], gltf["animations"][0]
    hb = gltf["extensions"]["VRMC_vrm_animation"]["humanoid"]["humanBones"]
    samplers = anim["samplers"]

    rot_out = {ch["target"]["node"]: samplers[ch["sampler"]]["output"]
               for ch in anim["channels"] if ch["target"]["path"] == "rotation"}
    trans_out = None
    hips_node = (hb.get("hips") or {}).get("node")
    for ch in anim["channels"]:
        if ch["target"]["path"] == "translation" and ch["target"]["node"] == hips_node:
            trans_out = samplers[ch["sampler"]]["output"]

    written = 0

    def _store(acc_idx, values):
        acc = gltf["accessors"][acc_idx]
        bv = gltf["bufferViews"][acc["bufferView"]]
        start = bv.get("byteOffset", 0)
        if acc["count"] != len(values):
            raise ValueError(
                f"frame count differs: template {acc['count']} vs data {len(values)}")
        data = np.ascontiguousarray(values, dtype=np.float32)
        if data.nbytes != bv["byteLength"]:
            raise ValueError("bufferView length mismatch")
        blob[start:start + data.nbytes] = data.tobytes()
        acc["min"] = [float(v) for v in data.min(axis=0)]
        acc["max"] = [float(v) for v in data.max(axis=0)]

    for slot, bone in enumerate(bones):
        entry = hb.get(bone)
        if not entry or entry.get("node") is None:
            continue
        node = entry["node"]
        if node not in rot_out:
            continue
        rest = np.asarray(nodes[node].get("rotation", [0.0, 0.0, 0.0, 1.0]),
                          dtype=np.float64)
        q = np.asarray(basis[:, slot], dtype=np.float64)
        q = q / np.clip(np.linalg.norm(q, axis=-1, keepdims=True), 1e-8, None)
        channel = _qmul(np.tile(rest, (len(q), 1)), q)
        _store(rot_out[node], channel)
        written += 1

    if trans is not None and trans_out is not None:
        _store(trans_out, np.asarray(trans, dtype=np.float32))

    json_chunk = json.dumps(gltf, separators=(",", ":")).encode("utf-8")
    json_chunk += b" " * ((-len(json_chunk)) % 4)
    bin_chunk = bytes(blob)
    bin_chunk += b"\0" * ((-len(bin_chunk)) % 4)
    total = 12 + 8 + len(json_chunk) + 8 + len(bin_chunk)
    with open(out_path, "wb") as f:
        f.write(struct.pack("<4sII", b"glTF", 2, total))
        f.write(struct.pack("<I4s", len(json_chunk), b"JSON"))
        f.write(json_chunk)
        f.write(struct.pack("<I4s", len(bin_chunk), b"BIN\0"))
        f.write(bin_chunk)
    return written


def channel_deviation(a_path, b_path):
    """Max |component| difference between two vrma files' rotation channels."""
    ga, ba = _read_glb(a_path)
    gb, bb = _read_glb(b_path)
    worst = 0.0
    for gltf, blob, other in ((ga, ba, (gb, bb)),):
        anim = gltf["animations"][0]
        for ch in anim["channels"]:
            if ch["target"]["path"] != "rotation":
                continue
            out = anim["samplers"][ch["sampler"]]["output"]
            x = _accessor(gltf, blob, out).astype(np.float64)
            y = _accessor(other[0], other[1], out).astype(np.float64)
            sgn = np.sign(np.sum(x * y, axis=-1, keepdims=True))
            sgn[sgn == 0] = 1.0
            worst = max(worst, float(np.abs(x - y * sgn).max()))
    return worst


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--npz", required=True, help="preprocessed solo GT clip")
    ap.add_argument("--gt-turn", required=True,
                    help="existing GT turn id to copy face.json/audio.wav from")
    ap.add_argument("--id", required=True)
    ap.add_argument("--emotion", default="neutral")
    ap.add_argument("--gpu", type=int, default=0)
    ap.add_argument("--no-self-verify", action="store_true")
    ap.add_argument("--force", action="store_true")
    args = ap.parse_args()

    gt_dir = os.path.join(HERE, "outputs", "turns", args.gt_turn)
    for name in ("face.json", "audio.wav"):
        if not os.path.isfile(os.path.join(gt_dir, name)):
            raise FileNotFoundError(f"{gt_dir}/{name} — render the GT turn first")

    turn_dir = os.path.join(HERE, "outputs", "turns", args.id)
    if os.path.exists(turn_dir) and not args.force:
        raise FileExistsError(
            f"{turn_dir} already exists; pass --force or choose another id")

    started = time.time()
    with np.load(args.npz, allow_pickle=False) as data:
        q_gt = data["vrm_rotations"].astype(np.float32)
        trans = data["trans"].astype(np.float32)
        source_bones = [str(v) for v in data["bone_names"]]
        dataset_id = str(data["dataset_id"])
        split = str(data["split"])
        segment_type = str(data["segment_type"])
        speech_ratio = float(data["speech_ratio"])
    if q_gt.shape[1:] != (NB, 4):
        raise ValueError(f"expected (T,{NB},4) GT rotations, got {q_gt.shape}")
    if source_bones != BONES:
        raise ValueError(f"dataset bone order does not match the {NB}-bone contract")

    # 회전값만 교체할 템플릿 = GT 턴이 그대로 싣는 바로 그 베이크 산출물.
    baked = os.path.normpath(os.path.join(
        os.path.dirname(args.npz), "..", "..", "vrma", split,
        os.path.splitext(os.path.basename(args.npz))[0] + ".vrma"))
    if not os.path.isfile(baked):
        raise FileNotFoundError(f"baked vrma not found: {baked}")

    os.makedirs(turn_dir, exist_ok=True)
    body_path = os.path.join(turn_dir, "body.vrma")

    # 자기검증: GT basis를 같은 writer로 되쓰면 베이크 채널이 재현돼야 한다.
    # 재현되면 이후 A/B의 차이는 writer가 아니라 코덱뿐이다.
    self_verify = None
    if not args.no_self_verify:
        probe = os.path.join(turn_dir, "_writer_selfcheck.vrma")
        write_channels(baked, q_gt, trans, probe)
        self_verify = channel_deviation(baked, probe)
        os.remove(probe)
        if self_verify > 1e-5:
            raise RuntimeError(
                f"writer is not neutral: {self_verify:.3e} — A/B would be invalid")

    device = torch.device(f"cuda:{args.gpu}" if torch.cuda.is_available() else "cpu")
    models, stats, selection = load_codecs(device)

    t0 = time.time()
    q_rec, trans_rec = roundtrip(q_gt, trans, models, stats, device)
    codec_s = time.time() - t0

    # Same clip, same numbers as the ceiling gate — but for THIS clip only, so the
    # viewer impression and the measurement refer to exactly the same frames.
    per_part = {
        part: {
            "mean_angle_deg": float(angular_error_deg(q_gt[:, sl], q_rec[:, sl]).mean()),
            "p95_angle_deg": float(np.percentile(
                angular_error_deg(q_gt[:, sl], q_rec[:, sl]), 95)),
        }
        for part, (_, sl) in PARTS.items()
    }
    drift_cm = float(np.linalg.norm(trans_rec[-1] - trans[-1]) * 100.0)

    write_channels(baked, q_rec, trans_rec, body_path)

    # Byte-identical face/audio: the body is the only difference vs the GT turn.
    for name in ("face.json", "audio.wav"):
        shutil.copyfile(os.path.join(gt_dir, name), os.path.join(turn_dir, name))
    face_frames = json.load(open(os.path.join(turn_dir, "face.json")))["num_frames"]

    meta = {
        "id": args.id,
        "emotion": args.emotion,
        "fps": FPS,
        "face_frames": int(face_frames),
        "body_frames": int(len(q_rec)),
        "face": "face.json",
        "body": "body.vrma",
        "audio": "audio.wav",
        "body_source": "codec_roundtrip_continuous_no_generator",
        "body_writer": "baked_template_channel_rewrite",
        "writer_self_verify_max_component": self_verify,
        "body_vrma_template": os.path.abspath(baked),
        "run_id": RUN_ID,
        "decode_mode": "continuous",
        "codec_sha256": selection["sha256"],
        "compare_with": args.gt_turn,
        "clip_angle_error_deg": per_part,
        "root_translation": "codec velocity, integrated from the GT first frame",
        "translation_drift_cm": round(drift_cm, 2),
        "dataset_id": dataset_id,
        "dataset_split": split,
        "source_npz": os.path.abspath(args.npz),
        "segment_type": segment_type,
        "speech_ratio": speech_ratio,
        "timing": {
            "face_s": 0.0,
            "body_s": round(codec_s, 2),
            "total_s": round(time.time() - started, 2),
        },
    }
    with open(os.path.join(turn_dir, "meta.json"), "w") as f:
        json.dump(meta, f, ensure_ascii=False, indent=1)

    print(f"{args.id}: {len(q_rec)} frames [{split}], codec {codec_s:.2f}s")
    if self_verify is not None:
        print(f"  writer 자기검증 {self_verify:.2e} (중립)")
    for part, row in per_part.items():
        print(f"  {part:8s} mean {row['mean_angle_deg']:7.3f}°  "
              f"p95 {row['p95_angle_deg']:8.3f}°")
    print(f"  이동 종점 드리프트 {drift_cm:.2f}cm")
    print(turn_dir)


if __name__ == "__main__":
    main()