# Reallusion Character Creator / iClone FBX → Quaternius-skeleton GLB retargeter.
#
# Sibling of tools/mixamo_retarget.py — identical world-rotation-delta method (CC rigs
# rest in T-pose, verified on ActorBuild exports), adapted for the CC_Base_* skeleton:
#   - fixed "CC_Base_" naming, hips bone = CC_Base_Hip (parent of Pelvis+Waist)
#   - twist helpers (UpperarmTwist01...) are unmapped → ride their parent at rest offset
#   - game-optimized CC exports have 1-segment fingers (Index1, no Index2/3) — mapped
#     where present, absent segments ride the parent
#   - iClone FBX takes named "*TempMotion" (2 frames) mean the export had NO motion —
#     the take is skipped with a loud warning (re-export with Include Motion).
#
# Usage:
#   blender --background --python tools/cc_retarget.py -- \
#       --src "<dir-or-fbx>[,<more>…]" --out <out.glb> [--name <ClipName>] [--target <gltf>]
#       [--fingers track|static]   static = ignore captured finger articulation; fingers
#                                  hold the target rig's rest pose (open hand) and the
#                                  hand moves as one unit. Use for noisy AI-mocap fingers.
#       [--despike]                heal short AI-mocap glitches: frames whose world-space
#                                  acceleration spikes far above the clip's median get
#                                  re-interpolated from the good frames around them.
#                                  Only short windows are touched, so intentional sharp
#                                  moves survive. Extra important for ping-pong clips —
#                                  a glitch plays twice (forward + reversed).
#
# Looping note: seamless looping is owned by the ping-pong bake (animation repo
# pingpong_bake.py → <pack>_pp), which loops and returns to start by construction.
# The former --inplace / --loop-blend passes were removed 2026-07-17; resurrect from
# git history (commit 535719601) if a forward-loop clip ever needs them again.

import bpy, sys, os, re
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")

BONE_MAP = {
    "CC_Base_Hip": "pelvis",
    "CC_Base_Waist": "spine_01",
    "CC_Base_Spine01": "spine_02",
    "CC_Base_Spine02": "spine_03",
    "CC_Base_NeckTwist01": "neck_01",
    "CC_Base_Head": "Head",
}
for S, T in (("L", "l"), ("R", "r")):
    BONE_MAP[f"CC_Base_{S}_Clavicle"] = f"clavicle_{T}"
    BONE_MAP[f"CC_Base_{S}_Upperarm"] = f"upperarm_{T}"
    BONE_MAP[f"CC_Base_{S}_Forearm"] = f"lowerarm_{T}"
    BONE_MAP[f"CC_Base_{S}_Hand"] = f"hand_{T}"
    BONE_MAP[f"CC_Base_{S}_Thigh"] = f"thigh_{T}"
    BONE_MAP[f"CC_Base_{S}_Calf"] = f"calf_{T}"
    BONE_MAP[f"CC_Base_{S}_Foot"] = f"foot_{T}"
    BONE_MAP[f"CC_Base_{S}_ToeBase"] = f"ball_{T}"
    for i in (1, 2, 3):
        BONE_MAP[f"CC_Base_{S}_Thumb{i}"] = f"thumb_0{i}_{T}"
    # game-optimized CC rigs: one segment per finger
    BONE_MAP[f"CC_Base_{S}_Index1"] = f"index_01_{T}"
    BONE_MAP[f"CC_Base_{S}_Mid1"] = f"middle_01_{T}"
    BONE_MAP[f"CC_Base_{S}_Ring1"] = f"ring_01_{T}"
    BONE_MAP[f"CC_Base_{S}_Pinky1"] = f"pinky_01_{T}"


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 clip_name(path):
    base = os.path.splitext(os.path.basename(path))[0]
    base = base.replace(" - ", "_").replace(" ", "")
    return re.sub(r"[^A-Za-z0-9_\-]", "", base)


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 char_frame(rest, hips, head, upper_arm_l):
    u = (rest[head].translation - rest[hips].translation).normalized()
    l = (rest[upper_arm_l].to_quaternion() @ Vector((0, 1, 0))).normalized()
    f = l.cross(u).normalized()
    l = u.cross(f).normalized()
    return Matrix((l, u, f)).transposed().to_quaternion().normalized()


def action_fcurves(a):
    """All fcurves across Blender 5 slotted-action layers/strips (legacy fallback)."""
    if hasattr(a, "fcurves") and a.fcurves is not None:
        return list(a.fcurves)
    fcs = []
    for layer in getattr(a, "layers", []):
        for strip in layer.strips:
            for bag in strip.channelbags:
                fcs.extend(bag.fcurves)
    return fcs


def pick_action(src_arm, acts):
    """The armature's motion take: fcurves must reference CC_Base_ bones; longest wins.
    TempMotion 2-frame placeholders (export without motion) are rejected."""
    best, best_len = None, 0.0
    for a in acts:
        if not any("CC_Base_" in fc.data_path for fc in action_fcurves(a)):
            continue
        length = a.frame_range[1] - a.frame_range[0]
        if length > best_len:
            best, best_len = a, length
    if best is not None and best_len < 5 and "TempMotion" in best.name:
        raise RuntimeError(
            f"'{best.name}' is a {int(best_len)+1}-frame iClone placeholder — the FBX was "
            "exported WITHOUT motion. Re-export with Include Motion / Export Motion checked.")
    return best


def retarget_one(src_arm, tgt_arm, action, name, scn, unit):
    f0, f1 = int(action.frame_range[0]), int(action.frame_range[1])
    if src_arm.animation_data is None:
        src_arm.animation_data_create()
    src_arm.animation_data.action = action
    if getattr(action, "slots", None) and src_arm.animation_data.action_slot is None:
        src_arm.animation_data.action_slot = action.slots[0]

    inv_map = {v: k for k, v in BONE_MAP.items()}

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

    C = (char_frame(tgt_rest, "pelvis", "Head", "upperarm_l")
         @ char_frame(src_rest, "CC_Base_Hip", "CC_Base_Head", "CC_Base_L_Upperarm").inverted()).normalized()
    Cinv = C.inverted()

    h_src = (src_rest["CC_Base_Head"].translation - src_rest["CC_Base_Hip"].translation).length
    h_tgt = (tgt_rest["Head"].translation - tgt_rest["pelvis"].translation).length
    tscale = h_tgt / h_src if h_src > 1e-5 else 1.0

    new_act = bpy.data.actions.new(name)
    if tgt_arm.animation_data is None:
        tgt_arm.animation_data_create()
    tgt_arm.animation_data.action = new_act

    order = hierarchy_order(tgt_arm)
    tw_inv = tgt_arm.matrix_world.inverted()
    for pb in tgt_arm.pose.bones:
        pb.rotation_mode = "QUATERNION"

    src_hips_rest_t = src_rest["CC_Base_Hip"].translation

    for f in range(f0, f1 + 1):
        scn.frame_set(f)
        dg = bpy.context.evaluated_depsgraph_get()
        se = src_arm.evaluated_get(dg)
        spose = {pb.name: (src_arm.matrix_world @ pb.matrix) for pb in se.pose.bones}

        tgt_world = {}
        for tb in order:
            nt = tb.name
            rest_t = tgt_rest[nt]
            ns = inv_map.get(nt)
            if ns is None or ns not in spose:
                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
            q_delta = (spose[ns].to_quaternion().normalized() @ src_rest_q[ns].inverted()).normalized()
            q_delta = (C @ q_delta @ Cinv).normalized()
            q_world = (q_delta @ tgt_rest_q[nt]).normalized()
            if nt == "pelvis":
                d = C @ (spose[ns].translation - src_hips_rest_t) * tscale * unit
                t_world = rest_t.translation + d
            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=f)
            if tb.name == "pelvis" or tb.parent is None:
                pb.keyframe_insert("location", frame=f)

    tgt_arm.animation_data.action = None
    return new_act


DESPIKE_FACTOR = 200.0  # world accel must exceed FACTOR x the clip's median accel
DESPIKE_MAX_WIN = 12    # frames; longer regions are real motion, leave them alone
DESPIKE_PAD = 3


def _hermite_patch(kps, a, b):
    """Replace kps[a+1..b-1] with a cubic that matches value AND slope at both edges."""
    v = [k.co.y for k in kps]
    n = b - a
    m0 = (v[a] - v[a - 1]) * n
    m1 = (v[b + 1] - v[b]) * n
    for i in range(a + 1, b):
        t = (i - a) / n
        h00 = 2 * t ** 3 - 3 * t ** 2 + 1
        h10 = t ** 3 - 2 * t ** 2 + t
        h01 = -2 * t ** 3 + 3 * t ** 2
        h11 = t ** 3 - t ** 2
        kps[i].co.y = h00 * v[a] + h10 * m0 + h01 * v[b] + h11 * m1


def despike(act, tgt_arm, scn):
    """Heal short AI-mocap glitches. Detection runs in WORLD space (same metric as
    tools/anim_qc.py): a bone whose world acceleration towers over the clip's median
    for a few frames is a tracking error, not choreography. The flagged bone's local
    curves — and its parent's, where the glitch usually originates — are re-interpolated
    across the window with slope-matched cubics."""
    bones = sorted(set(BONE_MAP.values()))
    prev = tgt_arm.animation_data.action
    tgt_arm.animation_data.action = act
    if getattr(act, "slots", None) and tgt_arm.animation_data.action_slot is None:
        tgt_arm.animation_data.action_slot = act.slots[0]
    f0, f1 = (int(x) for x in act.frame_range)

    pos = {n: [] for n in bones}
    for f in range(f0, f1 + 1):
        scn.frame_set(f)
        dg = bpy.context.evaluated_depsgraph_get()
        te = tgt_arm.evaluated_get(dg)
        for n in bones:
            pos[n].append((tgt_arm.matrix_world @ te.pose.bones[n].matrix).translation.copy())
    tgt_arm.animation_data.action = prev

    acc = {n: [(p[i + 1] - 2 * p[i] + p[i - 1]).length for i in range(1, len(p) - 1)]
           for n, p in pos.items()}
    flat = sorted(x for arr in acc.values() for x in arr)
    thr = DESPIKE_FACTOR * (flat[len(flat) // 2] or 1e-9)

    by_bone = {}
    for fc in action_fcurves(act):
        if fc.data_path.startswith('pose.bones["'):
            by_bone.setdefault(fc.data_path.split('"')[1], []).append(fc)

    repaired = 0
    for n, arr in acc.items():
        bad = [i + 1 for i, a in enumerate(arr) if a > thr]  # -> pos/key index
        if not bad:
            continue
        spans, s, e = [], bad[0], bad[0]
        for i in bad[1:]:
            if i - e <= DESPIKE_PAD * 2:
                e = i
            else:
                spans.append((s, e))
                s = e = i
        spans.append((s, e))
        parent = tgt_arm.data.bones[n].parent
        chain = [n] + ([parent.name] if parent else [])
        for s, e in spans:
            a = max(1, s - DESPIKE_PAD)
            b = min(len(pos[n]) - 2, e + DESPIKE_PAD)
            if b - a > DESPIKE_MAX_WIN or b - a < 2:
                continue
            for bone in chain:
                for fc in by_bone.get(bone, []):
                    kps = sorted(fc.keyframe_points, key=lambda k: k.co.x)
                    _hermite_patch(kps, a, b)
                    fc.update()
            repaired += 1
            print(f"[cc]   --despike: {n} @ frame {f0 + s}"
                  f" ({(s) / 60:.2f}s) window {b - a}f via {'+'.join(chain)}")
    print(f"[cc]   --despike: repaired {repaired} glitch window(s), threshold {thr:.4f}")


def main():
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    args, flags, i = {}, set(), 0
    while i < len(argv):
        if argv[i] == "--despike":
            flags.add("despike")
            i += 1
        else:
            args[argv[i]] = argv[i + 1]
            i += 2
    srcs, out = args["--src"].split(","), args["--out"]
    target = args.get("--target", DEFAULT_TARGET_GLTF)
    forced_name = args.get("--name")
    fingers = args.get("--fingers", "track")
    if fingers not in ("track", "static"):
        raise SystemExit(f"--fingers must be 'track' or 'static', got '{fingers}'")
    if fingers == "static":
        # Unmapped bones ride their parent at rest offset, so dropping the finger
        # entries makes each hand move as one rigid unit in the rig's rest hand pose.
        dropped = [k for k, v in BONE_MAP.items()
                   if re.match(r"(thumb|index|middle|ring|pinky)_", v)]
        for k in dropped:
            del BONE_MAP[k]
        print(f"[cc] --fingers static: ignoring {len(dropped)} finger bone mappings")

    fbx_files = []
    for s in srcs:
        s = s.strip()
        if os.path.isdir(s):
            for root, _, files in os.walk(s):
                fbx_files += [os.path.join(root, x) for x in sorted(files) if x.lower().endswith(".fbx")]
        else:
            fbx_files.append(s)
    print(f"[cc] {len(fbx_files)} clips → {out}")

    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()
    scn = bpy.context.scene
    scn.render.fps = 60  # iClone exports at 60; glTF export resamples fine

    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"
    for o in [o for o in bpy.data.objects if o.type == "MESH"]:
        bpy.data.objects.remove(o, do_unlink=True)
    bpy.context.view_layer.update()
    apply_object_transform(tgt_arm)

    done = []
    for i, f in enumerate(fbx_files):
        name = forced_name if (forced_name and len(fbx_files) == 1) else clip_name(f)
        pre = set(bpy.data.objects)
        pre_actions = set(bpy.data.actions)
        bpy.ops.import_scene.fbx(filepath=f)
        new_objs = [o for o in bpy.data.objects if o not in pre]
        src_arm = next((o for o in new_objs
                        if o.type == "ARMATURE" and any("CC_Base_Hip" == b.name for b in o.data.bones)), None)
        new_acts = [a for a in bpy.data.actions if a not in pre_actions]
        if src_arm is None or not new_acts:
            print(f"[cc] SKIP (no CC armature/action): {f}")
        else:
            bpy.context.view_layer.update()
            unit = src_arm.scale.x  # capture BEFORE transform_apply (fcurves stay in cm)
            apply_object_transform(src_arm)
            act = pick_action(src_arm, new_acts)
            if act is None:
                print(f"[cc] SKIP (no armature take): {f}")
            else:
                baked = retarget_one(src_arm, tgt_arm, act, name, scn, unit)
                if "despike" in flags:
                    despike(baked, tgt_arm, scn)
                done.append(baked)
                print(f"[cc] {i + 1}/{len(fbx_files)} {name} "
                      f"({int(baked.frame_range[1])}f @60, from take '{act.name}')")
        for o in new_objs:
            bpy.data.objects.remove(o, do_unlink=True)
        for a in new_acts:
            bpy.data.actions.remove(a)

    if not done:
        raise RuntimeError("nothing retargeted — see messages above")
    for act in done:
        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"[cc] EXPORTED {len(done)} clips → {out}")


main()