# MediaPipe pose JSON → Quaternius-skeleton GLB animation (Blender headless).
#
# Sibling of tools/mixamo_retarget.py — same target rig, same NLA/export tail, but the
# source is not an armature: it's per-frame 3D world landmarks from the pose-estimation
# skill (extract_pose.py). Rotations are SOLVED, not transferred:
#   - limb bones (arms/legs/hands/feet): shortest-arc delta from the bone's REST world
#     direction to the landmark-derived direction, applied on top of the rest rotation —
#     preserves the rig's rest roll, so skinning stays sane (no twist control in v1).
#   - pelvis + spine chain: full-basis deltas built from the hip line and the hips→
#     shoulders up vector; spine_01..03 slerp from pelvis delta toward chest delta.
#   - pelvis HEIGHT: MediaPipe world landmarks are hip-centered (origin = mid-hips every
#     frame), so global translation is lost. We reconstruct the vertical: per-frame
#     pelvis-above-ankle extent vs the standing extent (95th pct) scales the rest height
#     — carries squats/crouches, which is most of a haka.
# Landmarks are box-smoothed over time before solving (mocap jitter).
#
# MediaPipe world axes: x=subject-image right, y=down, z=depth (smaller=nearer camera).
# Converted here to Blender Z-up with the subject facing -Y: b = (x, z, -y).
#
# Usage:
#   blender --background --python tools/mocap_retarget.py -- \
#       --pose <extract_pose.json> --out <out.glb> --name <ClipName> \
#       [--range lo:hi] [--smooth 5] [--target <gltf>] [--render-check <dir>]

import bpy, sys, os, json
import numpy as np
from mathutils import Matrix, Quaternion, Vector

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_TARGET_GLTF = os.path.join(
    REPO_ROOT, "assets", "quaternius", "base-characters",
    "Universal Base Characters[Standard]", "Base Characters", "Godot - UE",
    "Superhero_Male_FullBody.gltf")

# MediaPipe pose landmark indices
NOSE, L_SHO, R_SHO, L_ELB, R_ELB, L_WRI, R_WRI = 0, 11, 12, 13, 14, 15, 16
L_PINK, R_PINK, L_IDX, R_IDX = 17, 18, 19, 20
L_HIP, R_HIP, L_KNE, R_KNE, L_ANK, R_ANK, L_FT, R_FT = 23, 24, 25, 26, 27, 28, 31, 32

# target bone -> (landmark_from, landmark_to): desired bone direction (bone +Y points from→to)
AIM_BONES = {
    "upperarm_l": (L_SHO, L_ELB), "lowerarm_l": (L_ELB, L_WRI),
    "upperarm_r": (R_SHO, R_ELB), "lowerarm_r": (R_ELB, R_WRI),
    "thigh_l": (L_HIP, L_KNE), "calf_l": (L_KNE, L_ANK), "foot_l": (L_ANK, L_FT),
    "thigh_r": (R_HIP, R_KNE), "calf_r": (R_KNE, R_ANK), "foot_r": (R_ANK, R_FT),
}
HAND_BONES = {"hand_l": (L_WRI, L_IDX, L_PINK), "hand_r": (R_WRI, R_IDX, R_PINK)}
SPINE_WEIGHTS = {"spine_01": 0.35, "spine_02": 0.7, "spine_03": 1.0,
                 "neck_01": 1.0, "Head": 1.0}


def apply_object_transform(obj):
    bpy.ops.object.select_all(action="DESELECT")
    obj.select_set(True)
    bpy.context.view_layer.objects.active = obj
    bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    bpy.context.view_layer.update()


def hierarchy_order(arm):
    out, stack = [], [b for b in arm.data.bones if b.parent is None]
    while stack:
        b = stack.pop(0)
        out.append(b)
        stack.extend(b.children)
    return out


def torso_quat(pts):
    """Orthonormal torso basis from landmarks (Blender axes): columns = (subject-left,
    up, forward). Subject-left = l_hip - r_hip; up = mid-shoulders - mid-hips."""
    left = Vector(pts[L_HIP] - pts[R_HIP]).normalized()
    up = Vector((pts[L_SHO] + pts[R_SHO]) / 2 - (pts[L_HIP] + pts[R_HIP]) / 2).normalized()
    fwd = left.cross(up).normalized()
    left = up.cross(fwd).normalized()
    return Matrix((left, up, fwd)).transposed().to_quaternion().normalized()


def chest_quat(pts):
    left = Vector(pts[L_SHO] - pts[R_SHO]).normalized()
    up = Vector((pts[L_SHO] + pts[R_SHO]) / 2 - (pts[L_HIP] + pts[R_HIP]) / 2).normalized()
    fwd = left.cross(up).normalized()
    left = up.cross(fwd).normalized()
    return Matrix((left, up, fwd)).transposed().to_quaternion().normalized()


def main():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    args = dict(zip(argv[::2], argv[1::2]))
    pose_path, out, clip = args["--pose"], args["--out"], args.get("--name", "Mocap01")
    target = args.get("--target", DEFAULT_TARGET_GLTF)
    smooth = int(args.get("--smooth", 5))
    render_dir = args.get("--render-check")
    lo, hi = 0.0, 1e9
    if "--range" in args:
        a, _, b = args["--range"].partition(":")
        lo, hi = float(a or 0), float(b) if b else 1e9

    doc = json.load(open(pose_path))
    frames = [f for f in doc["frames"]
              if f.get("detected") and "world_landmarks" in f and lo <= f["t"] <= hi]
    if len(frames) < 4:
        raise RuntimeError(f"only {len(frames)} usable frames in range {lo}:{hi}")
    ts = np.array([f["t"] for f in frames]); ts -= ts[0]
    # MediaPipe (x right, y down, z depth) → Blender (x right, y=z_mp, z=-y_mp)
    raw = np.array([[[p["x"], p["z"], -p["y"]] for p in f["world_landmarks"]] for f in frames])
    if smooth > 1:  # box smooth over time, per landmark/axis
        k = np.ones(smooth) / smooth
        pad = smooth // 2
        padded = np.concatenate([raw[:1].repeat(pad, 0), raw, raw[-1:].repeat(pad, 0)])
        raw = np.apply_along_axis(lambda v: np.convolve(v, k, mode="valid"), 0, padded)[:len(ts)]
    print(f"[mocap] {len(frames)} frames, {ts[-1]:.1f}s, smooth={smooth}")

    # ── target rig ──
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()
    scn = bpy.context.scene
    scn.render.fps = 30
    bpy.ops.import_scene.gltf(filepath=target)
    tgt_arm = next(o for o in bpy.data.objects if o.type == "ARMATURE")
    tgt_arm.name = "Armature"
    bpy.context.view_layer.update()
    apply_object_transform(tgt_arm)

    tgt_rest = {b.name: (tgt_arm.matrix_world @ b.matrix_local) for b in tgt_arm.data.bones}
    tgt_rest_q = {n: m.to_quaternion().normalized() for n, m in tgt_rest.items()}

    # Char-frame alignment, mixamo_retarget style: A maps landmark-space vectors and
    # rotation deltas into target world. Landmark space after conversion: subject-left=+X,
    # up=+Z, facing=-Y (NOT identity — it carries a 90° X rotation vs standard basis).
    up_t = (tgt_rest["Head"].translation - tgt_rest["pelvis"].translation).normalized()
    left_t = (tgt_rest_q["upperarm_l"] @ Vector((0, 1, 0))).normalized()  # bone Y = along bone
    fwd_t = left_t.cross(up_t).normalized()
    left_t = up_t.cross(fwd_t).normalized()
    Ct = Matrix((left_t, up_t, fwd_t)).transposed().to_quaternion().normalized()
    Cs = Matrix((Vector((1, 0, 0)), Vector((0, 0, 1)), Vector((0, -1, 0)))).transposed() \
        .to_quaternion().normalized()
    A = (Ct @ Cs.inverted()).normalized()
    Ainv = A.inverted()
    Cs_inv = Cs.inverted()

    # Size + pelvis-height model
    h_tgt = (tgt_rest["Head"].translation - tgt_rest["pelvis"].translation).length
    pelvis_up = raw[:, [L_ANK, R_ANK], 2].min(axis=1) * -1.0        # pelvis height above lowest ankle
    standing_h = float(np.percentile(pelvis_up, 95))
    rest_pelvis = tgt_rest["pelvis"].translation.copy()

    act = bpy.data.actions.new(clip)
    if tgt_arm.animation_data is None:
        tgt_arm.animation_data_create()
    tgt_arm.animation_data.action = act
    order = hierarchy_order(tgt_arm)
    tw_inv = tgt_arm.matrix_world.inverted()
    for pb in tgt_arm.pose.bones:
        pb.rotation_mode = "QUATERNION"

    check = []
    for i in range(len(ts)):
        pts = raw[i]
        frame = int(round(ts[i] * 30)) + 1

        # delta from landmark rest basis (Cs), conjugated into target world axes
        q_pelvis_d = (A @ (torso_quat(pts) @ Cs_inv) @ Ainv).normalized()
        q_chest_d = (A @ (chest_quat(pts) @ Cs_inv) @ Ainv).normalized()

        tgt_world = {}
        for tb in order:
            nt = tb.name
            rest_t = tgt_rest[nt]
            q_world = None
            if nt == "pelvis":
                q_world = (q_pelvis_d @ tgt_rest_q[nt]).normalized()
            elif nt in SPINE_WEIGHTS:
                qd = Quaternion(q_pelvis_d).slerp(q_chest_d, SPINE_WEIGHTS[nt])
                q_world = (qd @ tgt_rest_q[nt]).normalized()
            elif nt in AIM_BONES:
                a, b = AIM_BONES[nt]
                d_lm = Vector(pts[b] - pts[a])
                if d_lm.length > 1e-6:
                    d_tgt = A @ d_lm
                    rest_dir = tgt_rest_q[nt] @ Vector((0, 1, 0))
                    q_delta = rest_dir.rotation_difference(d_tgt.normalized())
                    q_world = (q_delta @ tgt_rest_q[nt]).normalized()
            elif nt in HAND_BONES:
                w, ix, pk = HAND_BONES[nt]
                d_lm = Vector((pts[ix] + pts[pk]) / 2 - pts[w])
                if d_lm.length > 1e-6:
                    d_tgt = A @ d_lm
                    rest_dir = tgt_rest_q[nt] @ Vector((0, 1, 0))
                    q_delta = rest_dir.rotation_difference(d_tgt.normalized())
                    q_world = (q_delta @ tgt_rest_q[nt]).normalized()

            if q_world is None:  # unmapped: ride parent rigidly at rest offset
                if tb.parent is None:
                    tgt_world[nt] = rest_t
                else:
                    rel = tgt_rest[tb.parent.name].inverted() @ rest_t
                    tgt_world[nt] = tgt_world[tb.parent.name] @ rel
                continue

            if nt == "pelvis":
                ratio = float(np.clip(pelvis_up[i] / max(standing_h, 1e-5), 0.35, 1.15))
                t_world = Vector((rest_pelvis.x, rest_pelvis.y, rest_pelvis.z * ratio))
            else:
                rel = tgt_rest[tb.parent.name].inverted() @ rest_t
                t_world = (tgt_world[tb.parent.name] @ rel).translation
            tgt_world[nt] = Matrix.LocRotScale(t_world, q_world, Vector((1, 1, 1)))

        for tb in order:
            pb = tgt_arm.pose.bones[tb.name]
            m_arm = tw_inv @ tgt_world[tb.name]
            if tb.parent:
                m_parent_arm = tw_inv @ tgt_world[tb.parent.name]
                rel_rest = tb.parent.matrix_local.inverted() @ tb.matrix_local
                basis = rel_rest.inverted() @ (m_parent_arm.inverted() @ m_arm)
            else:
                basis = tb.matrix_local.inverted() @ m_arm
            pb.matrix_basis = basis
            pb.keyframe_insert("rotation_quaternion", frame=frame)
            if tb.name == "pelvis" or tb.parent is None:
                pb.keyframe_insert("location", frame=frame)

        # numeric sanity: solved elbow angle vs landmark elbow angle (should match closely)
        if i % max(1, len(ts) // 6) == 0:
            def ang(u, v):
                u, v = u.normalized(), v.normalized()
                return np.degrees(np.arccos(np.clip(u.dot(v), -1, 1)))
            lm = ang(Vector(pts[L_SHO] - pts[L_ELB]), Vector(pts[L_WRI] - pts[L_ELB]))
            ua = (tgt_world["upperarm_l"].to_quaternion() @ Vector((0, 1, 0)))
            la = (tgt_world["lowerarm_l"].to_quaternion() @ Vector((0, 1, 0)))
            rig = ang(-ua, la)
            check.append((ts[i], lm, rig))

    print("[mocap] elbow-angle sanity (t, landmark°, rig°):")
    for t, a, b in check:
        print(f"    t={t:5.1f}  lm={a:6.1f}  rig={b:6.1f}  d={abs(a-b):5.1f}")

    # optional visual check renders (with mesh, before it's dropped)
    if render_dir:
        os.makedirs(render_dir, exist_ok=True)
        cam_data = bpy.data.cameras.new("ChkCam")
        cam = bpy.data.objects.new("ChkCam", cam_data)
        scn.collection.objects.link(cam)
        cam.location = (0, -3.2, 1.0)
        cam.rotation_euler = (np.radians(87), 0, 0)
        scn.camera = cam
        sun = bpy.data.objects.new("Sun", bpy.data.lights.new("Sun", "SUN"))
        scn.collection.objects.link(sun)
        sun.rotation_euler = (np.radians(45), 0, np.radians(30))
        scn.render.resolution_x, scn.render.resolution_y = 480, 640
        for i in range(0, len(ts), max(1, len(ts) // 6)):
            scn.frame_set(int(round(ts[i] * 30)) + 1)
            scn.render.filepath = os.path.join(render_dir, f"chk_{ts[i]:05.1f}s.png")
            bpy.ops.render.render(write_still=True)
        print(f"[mocap] check renders → {render_dir}")

    # animations-only GLB (game binds clips to its own body)
    for o in [o for o in bpy.data.objects if o.type == "MESH"]:
        bpy.data.objects.remove(o, do_unlink=True)
    tgt_arm.animation_data.action = None
    tr = tgt_arm.animation_data.nla_tracks.new()
    tr.name = act.name
    tr.strips.new(act.name, 1, act)
    bpy.ops.object.select_all(action="DESELECT")
    tgt_arm.select_set(True)
    bpy.context.view_layer.objects.active = tgt_arm
    bpy.ops.export_scene.gltf(
        filepath=out, use_selection=True,
        export_animations=True, export_animation_mode="NLA_TRACKS",
        export_force_sampling=True, export_optimize_animation_size=False)
    print(f"[mocap] EXPORTED '{clip}' ({len(ts)} keyed frames, {ts[-1]:.1f}s) → {out}")


main()