"""Export a reusable Mixamo rest pose, or fit it to a humanoid mesh.

Blender:
  blender -b --python Tools/model_tools/bl_fit_mixamo.py -- template <src.fbx> <dst.fbx>
  blender -b --python Tools/model_tools/bl_fit_mixamo.py -- fit <template.fbx> <mesh.fbx> <out.glb> [palm|full]
  blender -b --python Tools/model_tools/bl_fit_mixamo.py -- remap <fitted.glb> <tokenrig.glb> <out.glb>
"""
import json
import sys
from pathlib import Path

import bpy
import numpy as np
from math import pi
from mathutils import Matrix, Vector

BODY = [
    "mixamorig:Hips", "mixamorig:Spine", "mixamorig:Spine1", "mixamorig:Spine2",
    "mixamorig:Neck", "mixamorig:Head", "mixamorig:HeadTop_End",
    "mixamorig:LeftShoulder", "mixamorig:LeftArm", "mixamorig:LeftForeArm", "mixamorig:LeftHand",
    "mixamorig:RightShoulder", "mixamorig:RightArm", "mixamorig:RightForeArm", "mixamorig:RightHand",
    "mixamorig:LeftUpLeg", "mixamorig:LeftLeg", "mixamorig:LeftFoot", "mixamorig:LeftToeBase", "mixamorig:LeftToe_End",
    "mixamorig:RightUpLeg", "mixamorig:RightLeg", "mixamorig:RightFoot", "mixamorig:RightToeBase", "mixamorig:RightToe_End",
]


def argv():
    return sys.argv[sys.argv.index("--") + 1:]


def mesh_objects():
    return [o for o in bpy.context.scene.objects if o.type == "MESH"]


def armature():
    arms = [o for o in bpy.context.scene.objects if o.type == "ARMATURE"]
    if len(arms) != 1:
        raise SystemExit(f"expected one armature, found {len(arms)}")
    return arms[0]


def apply_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)


def bone_heads(arm):
    return {b.name: arm.matrix_world @ b.head_local for b in arm.data.bones}


def points_of(obj):
    if obj.type == "MESH":
        return np.array([obj.matrix_world @ v.co for v in obj.data.vertices], dtype=np.float64)
    return np.array([obj.matrix_world @ b.head_local for b in obj.data.bones], dtype=np.float64)


def rotate_axis_to_y(obj, axis):
    if axis == 1:
        return
    rotation = Matrix.Rotation(-pi / 2, 4, "X") if axis == 2 else Matrix.Rotation(pi / 2, 4, "Z")
    obj.matrix_world = rotation @ obj.matrix_world
    apply_transform(obj)


def put_height_on_y(obj):
    if obj.type == "ARMATURE":
        heads = bone_heads(obj)
        delta = heads["mixamorig:Head"] - heads["mixamorig:Hips"]
        axis = int(np.argmax(np.abs([delta.x, delta.y, delta.z])))
    else:
        span = points_of(obj).max(0) - points_of(obj).min(0)
        order = np.argsort(span)
        axis = int(order[-1] if span[order[-1]] > span[order[-2]] * 1.25 else order[-2])
    rotate_axis_to_y(obj, axis)


def face_negative_z(arm):
    heads = bone_heads(arm)
    forward = heads["mixamorig:LeftToeBase"] - heads["mixamorig:LeftFoot"]
    if abs(forward.x) > abs(forward.z):
        arm.matrix_world = Matrix.Rotation(pi / 2 if forward.x > 0 else -pi / 2, 4, "Y") @ arm.matrix_world
        apply_transform(arm)
        heads = bone_heads(arm)
        forward = heads["mixamorig:LeftToeBase"] - heads["mixamorig:LeftFoot"]
    if forward.z > 0:
        arm.matrix_world = Matrix.Rotation(pi, 4, "Y") @ arm.matrix_world
        apply_transform(arm)


def clear_motion(arm):
    if arm.animation_data:
        arm.animation_data_clear()
    arm.data.pose_position = "REST"


def export_fbx(arm, path):
    bpy.ops.object.select_all(action="DESELECT")
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.export_scene.fbx(
        filepath=str(path), use_selection=True, object_types={"ARMATURE"},
        add_leaf_bones=False, bake_anim=False,
        axis_forward="-Z", axis_up="Y", bake_space_transform=False,
    )


def template(src, dst):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.fbx(filepath=str(src))
    for obj in list(mesh_objects()):
        bpy.data.objects.remove(obj, do_unlink=True)
    arm = armature()
    clear_motion(arm)
    put_height_on_y(arm)
    face_negative_z(arm)
    heads = bone_heads(arm)
    feet = [heads[n] for n in heads if n.endswith("Foot") or n.endswith("ToeBase") or n.endswith("Toe_End")]
    min_y = min(p.y for p in feet)
    center_x = (heads["mixamorig:LeftUpLeg"].x + heads["mixamorig:RightUpLeg"].x) * 0.5
    center_z = heads["mixamorig:Hips"].z
    arm.location = Vector((-center_x, -min_y, -center_z))
    apply_transform(arm)
    dst = Path(dst)
    dst.parent.mkdir(parents=True, exist_ok=True)
    export_fbx(arm, dst)
    heads = bone_heads(arm)
    top = max(p.y for p in heads.values())
    info = {
        "bones": len(arm.data.bones),
        "names": [b.name for b in arm.data.bones],
        "height": round(top, 4),
        "up": "Y",
        "front": "-Z",
        "left": "-X",
        "source": str(src),
    }
    dst.with_suffix(".json").write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")
    print("TEMPLATE", json.dumps({k: info[k] for k in ("bones", "height", "up", "front")}))


def slice_centers(pts, y0, y1):
    sl = pts[(pts[:, 1] >= y0) & (pts[:, 1] < y1)]
    return sl


def landmarks(pts):
    lo, hi = pts.min(0), pts.max(0)
    height = hi[1] - lo[1]
    bands = np.linspace(lo[1], hi[1], 36)
    spans = []
    for a, b in zip(bands[:-1], bands[1:]):
        sl = slice_centers(pts, a, b)
        spans.append(0 if len(sl) < 8 else sl[:, 0].max() - sl[:, 0].min())
    arm_i = int(np.argmax(spans))
    arm_y = (bands[arm_i] + bands[arm_i + 1]) * 0.5
    arm = slice_centers(pts, arm_y - height * 0.04, arm_y + height * 0.04)
    reach = max(abs(arm[:, 0].min()), abs(arm[:, 0].max()))
    left_hand = arm[arm[:, 0] < -reach * 0.55].mean(0)
    right_hand = arm[arm[:, 0] > reach * 0.55].mean(0)
    torso = arm[np.abs(arm[:, 0]) < reach * 0.38]
    if len(torso) < 8:
        torso = arm
    shoulders = np.array([torso[:, 0].min(), arm_y, torso[:, 2].mean()])
    feet = pts[pts[:, 1] < lo[1] + height * 0.12]
    left_foot = feet[feet[:, 0] < 0].mean(0)
    right_foot = feet[feet[:, 0] > 0].mean(0)
    crotch_y = arm_y
    for a, b in zip(bands[:arm_i][::-1], bands[1:arm_i + 1][::-1]):
        sl = slice_centers(pts, a, b)
        if len(sl) < 8:
            continue
        xs = np.sort(sl[:, 0])
        if np.diff(xs).max() > reach * 0.18:
            crotch_y = (a + b) * 0.5
            break
    hip_y = min(arm_y - height * 0.05, crotch_y + height * 0.06)
    head = pts[pts[:, 1] > lo[1] + height * 0.78]
    return {
        "height": height,
        "arm_y": arm_y,
        "hip_y": hip_y,
        "head": head.mean(0),
        "top": hi[1],
        "left_hand": left_hand,
        "right_hand": right_hand,
        "left_shoulder_x": float(torso[:, 0].min()),
        "right_shoulder_x": float(torso[:, 0].max()),
        "left_foot": left_foot,
        "right_foot": right_foot,
        "depth": head[:, 2].mean() if False else float(head[:, 2].mean()),
    }


def chain_fraction(arm, names):
    ys = [(arm.matrix_world @ arm.data.bones[n].head_local).y for n in names if n in arm.data.bones]
    span = ys[-1] - ys[0]
    if abs(span) < 1e-6:
        return [0.0 for _ in ys]
    return [(y - ys[0]) / span for y in ys]


def move_bones(arm, marks):
    bpy.ops.object.select_all(action="DESELECT")
    arm.select_set(True)
    bpy.context.view_layer.objects.active = arm
    torso_names = [n for n in ("mixamorig:Hips", "mixamorig:Spine", "mixamorig:Spine1", "mixamorig:Spine2", "mixamorig:Neck", "mixamorig:Head", "mixamorig:HeadTop_End") if n in arm.data.bones]
    torso_fr = chain_fraction(arm, torso_names)
    side_fr = {}
    for side in ("Left", "Right"):
        arm_names = [f"mixamorig:{side}{p}" for p in ("Shoulder", "Arm", "ForeArm", "Hand") if f"mixamorig:{side}{p}" in arm.data.bones]
        leg_names = [n for n in (f"mixamorig:{side}UpLeg", f"mixamorig:{side}Leg", f"mixamorig:{side}Foot") if n in arm.data.bones]
        side_fr[side] = (arm_names, chain_fraction(arm, arm_names), leg_names, chain_fraction(arm, leg_names))
    bpy.ops.object.mode_set(mode="EDIT")
    bones = arm.data.edit_bones
    rolls = {b.name: b.roll for b in bones}
    present = {b.name: b for b in bones}

    def put(name, head, tail):
        bone = present[name]
        bone.use_connect = False
        bone.head = Vector(head)
        bone.tail = Vector(tail)
        bone.roll = rolls[name]

    def between(a, b, t):
        return a * (1 - t) + b * t

    fr = torso_fr
    hip = np.array([0.0, marks["hip_y"], marks["depth"]])
    crown = np.array([marks["head"][0] * 0.2, marks["top"], marks["head"][2]])
    placed = {}
    for name, t in zip(torso_names, fr):
        placed[name] = between(hip, crown, t)
    for i, name in enumerate(torso_names):
        nxt = placed[torso_names[i + 1]] if i + 1 < len(torso_names) else placed[name] + np.array([0, marks["height"] * 0.03, 0])
        put(name, placed[name], nxt)

    for side, sign, hand, foot, shoulder_x in (
        ("Left", -1, marks["left_hand"], marks["left_foot"], marks["left_shoulder_x"]),
        ("Right", 1, marks["right_hand"], marks["right_foot"], marks["right_shoulder_x"]),
    ):
        shoulder = np.array([shoulder_x, marks["arm_y"], marks["depth"]])
        hand = np.array(hand)
        arm_names, afr, leg_core, lfr = side_fr[side]
        joints = [between(shoulder, hand, t) for t in afr]
        for i, name in enumerate(arm_names):
            nxt = joints[i + 1] if i + 1 < len(joints) else joints[i] + (hand - shoulder) * 0.08
            put(name, joints[i], nxt)
        hip_j = np.array([foot[0] * 0.72, marks["hip_y"], marks["depth"]])
        foot_j = np.array([foot[0], max(foot[1], 0.0), foot[2]])
        leg_names = leg_core + [n for n in (f"mixamorig:{side}ToeBase", f"mixamorig:{side}Toe_End") if n in present]
        leg = [between(hip_j, foot_j, t) for t in lfr]
        toe = foot_j + np.array([0.0, 0.0, -abs(marks["height"]) * 0.06])
        toe_end = toe + np.array([0.0, 0.0, -abs(marks["height"]) * 0.035])
        extra = {f"mixamorig:{side}ToeBase": toe, f"mixamorig:{side}Toe_End": toe_end}
        seq = leg + [extra[n] for n in leg_names[3:]]
        for i, name in enumerate(leg_names):
            nxt = seq[i + 1] if i + 1 < len(seq) else seq[i] + np.array([0, 0, -0.02])
            put(name, seq[i], nxt)
    bpy.ops.object.mode_set(mode="OBJECT")


def drop_fingers(arm):
    bpy.ops.object.mode_set(mode="EDIT")
    keys = ("Thumb", "Index", "Middle", "Ring", "Pinky")
    guard = 0
    while guard < 8:
        guard += 1
        bones = list(arm.data.edit_bones)
        leaves = [b for b in bones if any(k in b.name for k in keys) and not any(c.parent == b for c in bones)]
        if not leaves:
            break
        for bone in leaves:
            arm.data.edit_bones.remove(bone)
    bpy.ops.object.mode_set(mode="OBJECT")


def fit(template_path, mesh_path, out_path, fingers):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.fbx(filepath=str(mesh_path))
    mesh = mesh_objects()[0]
    put_height_on_y(mesh)
    before = set(bpy.data.objects)
    bpy.ops.import_scene.fbx(filepath=str(template_path))
    arm = next(o for o in bpy.context.scene.objects if o not in before and o.type == "ARMATURE")
    clear_motion(arm)
    put_height_on_y(arm)
    face_negative_z(arm)
    if fingers == "palm":
        drop_fingers(arm)
    marks = landmarks(points_of(mesh))
    move_bones(arm, marks)
    heads = bone_heads(arm)
    print("MARKS", json.dumps({k: (round(float(v), 3) if np.ndim(v) == 0 else [round(float(x), 3) for x in v]) for k, v in marks.items() if k != "head"}))
    for name in ("mixamorig:Hips", "mixamorig:Head", "mixamorig:LeftHand", "mixamorig:LeftFoot", "mixamorig:Spine2"):
        p = heads[name]
        print(name, round(p.x, 3), round(p.y, 3), round(p.z, 3))
    mesh.parent = arm
    arm.matrix_world = Matrix.Rotation(pi / 2, 4, "X") @ arm.matrix_world
    apply_transform(arm)
    out = Path(out_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    bpy.ops.export_scene.gltf(filepath=str(out), export_format="GLB", export_yup=True, export_skins=True)
    names = [b.name for b in arm.data.bones]
    print("FIT", json.dumps({"bones": len(names), "fingers": fingers, "out": str(out)}))


def remap_names(source_glb, rigged_glb, out_glb):
    """TokenRig --use_skeleton renames joints to bone_* and can reverse a limb chain.

    Match each output joint to the fitted Mixamo bone by head position, then copy
    that bone's parent, head, tail and roll. Vertex groups follow the new names.
    """
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=str(source_glb))
    src = armature()
    bpy.context.view_layer.objects.active = src
    bpy.ops.object.mode_set(mode="EDIT")
    src_pose = {}
    for bone in src.data.edit_bones:
        src_pose[bone.name] = {
            "parent": bone.parent.name if bone.parent else None,
            "head": src.matrix_world @ bone.head,
            "tail": src.matrix_world @ bone.tail,
            "roll": bone.roll,
        }
    bpy.ops.object.mode_set(mode="OBJECT")
    for obj in list(bpy.data.objects):
        bpy.data.objects.remove(obj, do_unlink=True)

    bpy.ops.import_scene.gltf(filepath=str(rigged_glb))
    arm = armature()
    mesh = mesh_objects()[0]
    used = set()
    mapping = {}
    for bone in arm.data.bones:
        head = arm.matrix_world @ bone.head_local
        name, dist = min(
            ((n, (head - pose["head"]).length) for n, pose in src_pose.items() if n not in used),
            key=lambda item: item[1],
        )
        used.add(name)
        mapping[bone.name] = (name, dist)
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode="EDIT")
    for old, (new, _dist) in mapping.items():
        arm.data.edit_bones[old].name = new
    for bone in arm.data.edit_bones:
        bone.use_connect = False
        bone.parent = None
    for bone in arm.data.edit_bones:
        pose = src_pose[bone.name]
        bone.head = arm.matrix_world.inverted() @ pose["head"]
        bone.tail = arm.matrix_world.inverted() @ pose["tail"]
        bone.roll = pose["roll"]
    for bone in arm.data.edit_bones:
        parent = src_pose[bone.name]["parent"]
        bone.parent = arm.data.edit_bones[parent] if parent else None
    bpy.ops.object.mode_set(mode="OBJECT")
    for group in mesh.vertex_groups:
        if group.name in mapping:
            group.name = mapping[group.name][0]
    worst = max(dist for _, dist in mapping.values())
    if worst > 0.05:
        raise SystemExit(f"bone heads moved, worst={worst:.4f}")
    out = Path(out_glb)
    out.parent.mkdir(parents=True, exist_ok=True)
    bpy.ops.export_scene.gltf(filepath=str(out), export_format="GLB", export_yup=True, export_skins=True)
    print("REMAP", json.dumps({"bones": len(mapping), "worst": round(worst, 4), "out": str(out)}))


def main():
    args = argv()
    if args[0] == "template":
        template(args[1], args[2])
    elif args[0] == "fit":
        fit(args[1], args[2], args[3], args[4] if len(args) > 4 else "palm")
    elif args[0] == "remap":
        remap_names(args[1], args[2], args[3])
    else:
        raise SystemExit("use template, fit, or remap")


if __name__ == "__main__":
    main()