"""
The player's body, animations (module of the import runner, see import_assets.py). 2026-10-12, model-independent since the third task.

Reads Config/DescentBody.json (the SAME file the game reads): the model's profile (mesh, animation folder, bone names) and the
animation table (situation role -> source clip on the UE5 mannequin). For the model named by the environment variable
DESCENT_BODY_MODEL (default: the file's "model") it RETARGETS every source clip onto the model's skeleton and BAKES one sequence
per role: <anim_folder>/A_Body_<Role>. The game loads by that name. Nothing is saved outside <anim_folder> and its sibling
Retarget folder: the user's packs are only read.

How: one IK Rig per skeletal mesh (unreal.IKRigController.apply_auto_generated_retarget_definition), one IK Retargeter per source
pack (default ops, chains mapped by name, target pose auto-aligned), unreal.IKRetargetBatchOperation.run_batch_retarget into
temporary TMP_ sequences, which are then copied to their role names. A model whose bones the engine cannot characterise by name
needs its IK Rig made by hand (see Docs/PLAYER_MODEL_SWAP.md).

Afterwards every baked clip is measured and the numbers the game needs are written next to the clips, on disk, as
<anim_folder>/BodyAnims.measured.json: length, and for walking and running loops the speed the clip moves at (the planted foot
against the body), how unevenly that foot moves while planted (the slide left after matching the play rate), the footfalls per
clip and the phase of the pelvis' up-and-down (the view bob follows it).

Re-run: skipped while the table, the profile and the baked sequences are unchanged; `-Force` bakes again. Sequences of roles that
left the table are deleted (they are ours).
"""
import json
import math
import os

import unreal

PREFIX = "A_Body_"
TMP = "TMP_"
SRC, TGT = unreal.RetargetSourceOrTarget.SOURCE, unreal.RetargetSourceOrTarget.TARGET
LOCO = ("Walk", "Jog", "Run")


def _fresh(ctx, folder, name, cls, factory):
    # A rig or retargeter of an earlier run is used again and set up anew below. (Deleting it first does not work unattended:
    # the retargeters and the baked sequences refer to it, and the editor then refuses without a dialog.)
    return ctx.get_or_create(folder, name, cls, factory)


def _make_rig(ctx, folder, name, mesh):
    rig = _fresh(ctx, folder, name, unreal.IKRigDefinition, unreal.IKRigDefinitionFactory())
    c = unreal.IKRigController.get_controller(rig)
    if not c.set_skeletal_mesh(mesh):
        raise RuntimeError("IK rig %s: mesh %s refused" % (name, mesh.get_name()))
    if not c.apply_auto_generated_retarget_definition():
        raise RuntimeError("IK rig %s: the engine could not characterise %s by its bone names; make the rig by hand" % (name, mesh.get_name()))
    chains = [str(ch.chain_name) for ch in c.get_retarget_chains()]
    ctx.log("rig %s on %s: root %s, %d chains" % (name, mesh.get_name(), c.get_retarget_root(), len(chains)))
    return rig, chains


def _source_mesh(ctx, seq_path):
    """A skeletal mesh of the sequence's own pack that sits on the sequence's own skeleton asset."""
    seq = ctx.eal.load_asset(seq_path)
    if seq is None:
        raise RuntimeError("source animation %s is missing" % seq_path)
    skel = seq.get_editor_property("skeleton")
    pack = "/".join(seq_path.split("/")[:3])
    found = []
    for data in unreal.AssetRegistryHelpers.get_asset_registry().get_assets_by_path(pack, recursive=True):
        if str(data.asset_class_path.asset_name) == "SkeletalMesh":
            mesh = data.get_asset()
            if mesh and mesh.get_editor_property("skeleton") == skel:
                found.append(mesh)
    if not found:
        raise RuntimeError("no skeletal mesh in %s on skeleton %s (needed to retarget %s)" % (pack, skel.get_path_name(), seq_path))
    found.sort(key=lambda m: (0 if "Manny" in m.get_name() else 1, 0 if "Simple" in m.get_name() else 1, m.get_name()))
    return seq, found[0]


def _in_root(seq, bone, t):
    """Location of a bone in the ROOT's frame (the root's own travel left out: what an in-place clip shows)."""
    path = [str(n) for n in unreal.AnimationLibrary.find_bone_path_to_root(seq, bone)]
    tr = unreal.AnimationLibrary.get_bone_pose_for_time(seq, path[0], t, False)
    for parent in path[1:-1]:
        tr = unreal.MathLibrary.compose_transforms(tr, unreal.AnimationLibrary.get_bone_pose_for_time(seq, parent, t, False))
    return tr.translation


def _measure(seq, role, bones):
    length = seq.get_play_length()
    out = {"length": round(length, 4)}
    if not role.startswith(LOCO):
        return out
    n = 96
    dt = length / n
    times = [dt * k for k in range(n)]
    pel = [_in_root(seq, bones["pelvis"], t) for t in times]
    # the pelvis' up-and-down as one harmonic of the clip: z ~ amp * cos(2 pi * steps * t / length - phase)
    best = (0.0, 1, 0.0)
    for h in range(1, 9):
        a = sum(pel[k].z * math.cos(2 * math.pi * h * k / n) for k in range(n))
        b = sum(pel[k].z * math.sin(2 * math.pi * h * k / n) for k in range(n))
        amp = 2.0 / n * math.hypot(a, b)
        if amp > best[0]:
            best = (amp, h, math.atan2(b, a))
    out.update({"bob_steps": best[1], "bob_phase": round(best[2], 3), "bob_amp_cm": round(best[0], 2)})
    # The planted foot moves backward against the body at the speed the clip walks at. Planted = the lower foot, near its lowest.
    feet = [[_in_root(seq, bones[f], t) for t in times] for f in ("foot_l", "foot_r")]
    floor = min(p.z for f in feet for p in f)
    speeds = []
    for k in range(n):
        f = feet[0] if feet[0][k].z <= feet[1][k].z else feet[1]
        if f[k].z <= floor + 2.5:
            nxt = f[(k + 1) % n]
            speeds.append(math.hypot(nxt.x - f[k].x, nxt.y - f[k].y) / dt)
    if speeds:
        speeds.sort()
        speed = speeds[len(speeds) // 2]
        out["speed_cm_s"] = round(speed, 1)
        out["slide_cm_s"] = round(sum(abs(s - speed) for s in speeds) / len(speeds), 1)   # what is left of skating at the matched rate
        out["stride_cm"] = round(speed * length / max(1, best[1]), 1)                     # ground covered per footfall
    root0 = unreal.AnimationLibrary.get_bone_pose_for_time(seq, "root", 0.0, False).translation
    root1 = unreal.AnimationLibrary.get_bone_pose_for_time(seq, "root", max(0.0, length - 1e-3), False).translation
    out["root_speed_cm_s"] = round(math.hypot(root1.x - root0.x, root1.y - root0.y) / length, 1)
    return out


def run(ctx):
    EAL = ctx.eal
    table_path = os.path.join(ctx.project_dir, "Config", "DescentBody.json")
    with open(table_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    model = os.environ.get("DESCENT_BODY_MODEL") or data["model"]
    if model not in data["profiles"]:
        raise RuntimeError("Config/DescentBody.json has no profile %r" % model)
    profile = data["profiles"][model]
    rows = data["animations"]["rows"]
    anim_dir = profile["anim_folder"].rstrip("/")
    rig_dir = anim_dir.rsplit("/", 1)[0] + "/Retarget"
    wanted = {PREFIX + r["role"]: r for r in rows}
    stamp = json.dumps({"mesh": profile["mesh"], "rows": [(r["role"], r["source"]) for r in rows]}, sort_keys=True)
    state = ctx.load_state("body_anims_" + model)
    have_all = all(EAL.does_asset_exist(anim_dir + "/" + n) for n in wanted)
    ctx.log("model %s: %d roles into %s" % (model, len(rows), anim_dir))

    if ctx.force or not have_all or state.get("stamp") != stamp:
        EAL.make_directory(anim_dir)
        EAL.make_directory(rig_dir)
        target = EAL.load_asset(profile["mesh"])
        if target is None:
            raise RuntimeError("the model's mesh %s is missing" % profile["mesh"])
        tgt_rig, tgt_chains = _make_rig(ctx, rig_dir, "IK_" + model, target)
        groups = {}   # source mesh -> the distinct source sequences retargeted through it
        for path in sorted(set(r["source"] for r in rows)):
            seq, mesh = _source_mesh(ctx, path)
            groups.setdefault(mesh.get_path_name(), (mesh, []))[1].append(seq)
        for mesh, seqs in groups.values():
            pack = mesh.get_path_name().split("/")[2]
            src_rig, _ = _make_rig(ctx, rig_dir, "IK_" + pack, mesh)
            rtg = _fresh(ctx, rig_dir, "RTG_%s_To_%s" % (pack, model), unreal.IKRetargeter, unreal.IKRetargetFactory())
            c = unreal.IKRetargeterController.get_controller(rtg)
            c.set_ik_rig(SRC, src_rig)
            c.set_ik_rig(TGT, tgt_rig)
            c.set_preview_mesh(SRC, mesh)
            c.set_preview_mesh(TGT, target)
            if c.get_num_retarget_ops() == 0:
                c.add_default_ops()
            c.assign_ik_rig_to_all_ops(SRC, src_rig)
            c.assign_ik_rig_to_all_ops(TGT, tgt_rig)
            c.auto_map_chains(unreal.AutoMapChainType.FUZZY, True)
            c.set_current_retarget_pose(c.create_retarget_pose("Aligned", TGT), TGT)
            c.auto_align_all_bones(TGT)
            unmapped = [n for n in tgt_chains if str(c.get_source_chain(n)) in ("None", "")]
            ctx.log("retargeter %s: %d of %d target chains mapped%s" % (rtg.get_name(), len(tgt_chains) - len(unmapped), len(tgt_chains),
                                                                        (", unmapped: " + ", ".join(unmapped)) if unmapped else ""))
            inp = unreal.IKRetargetBatchOperationInputs()
            for prop, value in (("assets_to_retarget", [unreal.AssetRegistryHelpers.create_asset_data(s) for s in seqs]), ("source_mesh", mesh),
                                ("target_mesh", target), ("ik_retarget_asset", rtg), ("prefix", TMP), ("target_path", anim_dir),
                                ("use_source_path", False), ("include_referenced_assets", False), ("overwrite_existing_files", True)):
                inp.set_editor_property(prop, value)
            made = unreal.IKRetargetBatchOperation.run_batch_retarget(inp)
            ctx.log("retargeted %d of %d sequences from %s" % (len(made), len(seqs), pack))
        # one sequence per role, by role name; then the temporaries go
        for name, row in wanted.items():
            tmp = anim_dir + "/" + TMP + row["source"].rsplit("/", 1)[1]
            if not EAL.does_asset_exist(tmp):
                ctx.fail("role %s: the retarget of %s made nothing" % (row["role"], row["source"]))
                continue
            if EAL.does_asset_exist(anim_dir + "/" + name):
                EAL.delete_asset(anim_dir + "/" + name)
            if EAL.duplicate_asset(tmp, anim_dir + "/" + name) is None:
                ctx.fail("role %s: could not copy %s to %s" % (row["role"], tmp, name))
        EAL.save_directory(rig_dir, only_if_is_dirty=False, recursive=True)
    else:
        ctx.log("table, profile and the %d baked sequences are unchanged: not retargeted again (-Force does)" % len(wanted))

    # ours that the table no longer names (and the temporaries, and the names of before the role table)
    for path in EAL.list_assets(anim_dir, recursive=False, include_folder=False):
        name = path.rsplit("/", 1)[1].split(".")[0]
        if name.startswith((TMP, "TB_")) or (name.startswith(PREFIX) and name not in wanted):
            EAL.delete_asset(anim_dir + "/" + name)
            ctx.log("removed %s" % name)

    measured = {}
    for name, row in wanted.items():
        seq = EAL.load_asset(anim_dir + "/" + name)
        if seq is None:
            ctx.fail("the baked sequence %s/%s is missing" % (anim_dir, name))
            continue
        lock = bool(row.get("root_lock", True))
        seq.set_editor_property("enable_root_motion", False)
        seq.set_editor_property("root_motion_root_lock", unreal.RootMotionRootLock.REF_POSE)
        seq.set_editor_property("force_root_lock", lock)
        measured[row["role"]] = _measure(seq, row["role"], profile["bones"])
        if "speed_cm_s" in measured[row["role"]]:
            ctx.log("%s %s" % (row["role"], json.dumps(measured[row["role"]], separators=(",", ":"))))
    EAL.save_directory(anim_dir, only_if_is_dirty=False, recursive=True)
    disk = os.path.join(unreal.Paths.convert_relative_path_to_full(unreal.Paths.project_content_dir()), anim_dir[len("/Game/"):], "BodyAnims.measured.json")
    with open(disk, "w", encoding="utf-8") as f:
        json.dump({"_about": "Written by Tools/ue/assets_14_body_anims.py; read by the game. Centimetres of the unscaled mesh.", "model": model, "roles": measured}, f, indent=1)
    ctx.log("measured %d roles -> %s" % (len(measured), disk))
    ctx.save_state("body_anims_" + model, {"stamp": stamp})