"""Fit a Kimodo SOMA rest pose onto a humanoid mesh without changing bone axes.

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

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

DROP = ("Thumb", "Index", "Middle", "Ring", "Pinky", "Jaw", "Eye")


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


def meshes():
    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):
    return np.array([obj.matrix_world @ v.co for v in obj.data.vertices], dtype=np.float64)


def rotate_up_to_z(obj):
    if obj.type == "ARMATURE":
        heads = bone_heads(obj)
        delta = heads["Head"] - heads["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)
        axis = int(np.argmax(span))
    if axis == 2:
        return
    rotation = Matrix.Rotation(pi / 2, 4, "X") if axis == 1 else Matrix.Rotation(-pi / 2, 4, "Y")
    obj.matrix_world = rotation @ obj.matrix_world
    apply_transform(obj)


def face_positive_y(arm):
    heads = bone_heads(arm)
    forward = heads["LeftToeBase"] - heads["LeftFoot"]
    flat = Vector((forward.x, forward.y, 0.0))
    if flat.length < 1e-6:
        return
    angle = flat.angle(Vector((0.0, 1.0, 0.0)), 0.0)
    cross = flat.x * 1.0 - flat.y * 0.0
    signed = angle if cross >= 0 else -angle
    arm.matrix_world = Matrix.Rotation(signed, 4, "Z") @ arm.matrix_world
    apply_transform(arm)
    heads = bone_heads(arm)
    if heads["LeftHand"].x > heads["RightHand"].x:
        arm.matrix_world = Matrix.Rotation(pi, 4, "Z") @ arm.matrix_world
        apply_transform(arm)


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


def plant(arm):
    heads = bone_heads(arm)
    feet = [heads[n] for n in heads if n.endswith("Foot") or "Toe" in n]
    hips = heads["Hips"]
    arm.location = Vector((-hips.x, -hips.y, -min(p.z for p in feet)))
    apply_transform(arm)


def extend_stubs(arm):
    bpy.ops.object.mode_set(mode="EDIT")
    for bone in arm.data.edit_bones:
        if (bone.tail - bone.head).length > 1e-4:
            continue
        direction = bone.parent.vector.normalized() if bone.parent and bone.parent.vector.length > 1e-6 else Vector((0.0, 1.0, 0.0))
        bone.tail = bone.head + direction * 0.015
    bpy.ops.object.mode_set(mode="OBJECT")


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


def template(bvh, dst):
    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_anim.bvh(filepath=str(bvh), global_scale=0.01, rotate_mode="QUATERNION")
    arm = armature()
    clear_motion(arm)
    rotate_up_to_z(arm)
    plant(arm)
    extend_stubs(arm)
    export_fbx(arm, dst)
    heads = bone_heads(arm)
    info = {
        "bones": len(arm.data.bones),
        "names": [b.name for b in arm.data.bones],
        "height": round(max(p.z for p in heads.values()) - min(p.z for p in heads.values()), 4),
        "up": "Z-in-blender",
        "front": "+Y-in-blender",
        "left": "-X",
        "source": str(bvh),
    }
    Path(dst).with_suffix(".json").write_text(json.dumps(info, ensure_ascii=False, indent=2), encoding="utf-8")
    hand = heads["LeftHand"]
    toe = heads["LeftToeBase"] - heads["LeftFoot"]
    print("TEMPLATE", json.dumps({"bones": info["bones"], "height": info["height"], "leftHandX": round(hand.x, 3), "toeY": round(toe.y, 3)}))


def landmarks(pts):
    lo, hi = pts.min(0), pts.max(0)
    height = hi[2] - lo[2]
    bands = np.linspace(lo[2] + height * 0.25, lo[2] + height * 0.8, 12)
    spans = []
    for a, b in zip(bands[:-1], bands[1:]):
        sl = pts[(pts[:, 2] >= a) & (pts[:, 2] < b)]
        spans.append(0 if len(sl) < 8 else sl[:, 0].max() - sl[:, 0].min())
    arm_i = int(np.argmax(spans))
    arm_z = (bands[arm_i] + bands[arm_i + 1]) * 0.5
    arm = pts[np.abs(pts[:, 2] - arm_z) < height * 0.04]
    reach = max(abs(arm[:, 0].min()), abs(arm[:, 0].max()))
    torso = arm[np.abs(arm[:, 0]) < reach * 0.38]
    if len(torso) < 8:
        torso = arm
    feet = pts[pts[:, 2] < lo[2] + height * 0.12]
    head = pts[pts[:, 2] > lo[2] + height * 0.78]
    hip_z = arm_z - height * 0.16
    return {
        "height": float(height),
        "arm_z": float(arm_z),
        "hip_z": float(hip_z),
        "head": head.mean(0),
        "top": float(hi[2]),
        "left_hand": arm[arm[:, 0] > reach * 0.55].mean(0),
        "right_hand": arm[arm[:, 0] < -reach * 0.55].mean(0),
        "left_foot": feet[feet[:, 0] > 0].mean(0),
        "right_foot": feet[feet[:, 0] < 0].mean(0),
        "depth": float(torso[:, 1].mean()),
    }


def subtree(bones, root):
    names = [root]
    pending = [root]
    while pending:
        parent = pending.pop()
        for bone in bones:
            if bone.parent and bone.parent.name == parent and bone.name not in names:
                names.append(bone.name)
                pending.append(bone.name)
    return names


def scale_from(bones, names, pivot, factor):
    for name in names:
        bone = bones[name]
        bone.head = pivot + (bone.head - pivot) * factor
        bone.tail = pivot + (bone.tail - pivot) * factor


def drop_extras(arm):
    bpy.ops.object.mode_set(mode="EDIT")
    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 DROP) 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 = max(meshes(), key=lambda item: len(item.data.vertices))
    for extra in list(meshes()):
        if extra != mesh:
            bpy.data.objects.remove(extra, do_unlink=True)
    rotate_up_to_z(mesh)
    mesh.matrix_world = Matrix.Rotation(pi, 4, "Z") @ mesh.matrix_world
    apply_transform(mesh)
    before = set(bpy.data.objects)
    if str(template_path).lower().endswith(".bvh"):
        bpy.ops.import_anim.bvh(filepath=str(template_path), global_scale=0.01, rotate_mode="QUATERNION")
    else:
        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)
    rotate_up_to_z(arm)
    plant(arm)
    if fingers == "palm":
        drop_extras(arm)
    extend_stubs(arm)
    marks = landmarks(points_of(mesh))
    bpy.context.view_layer.objects.active = arm
    bpy.ops.object.mode_set(mode="EDIT")
    bones = arm.data.edit_bones
    for bone in bones:
        bone.use_connect = False
    height = max(b.head.z for b in bones) - min(b.head.z for b in bones)
    scale_from(bones, [b.name for b in bones], Vector((0.0, 0.0, 0.0)), marks["height"] / height)
    hips = bones["Hips"].head.copy()
    shift = Vector((0.0, marks["depth"] - hips.y, marks["hip_z"] - hips.z))
    for bone in bones:
        bone.head += shift
        bone.tail += shift
    before_dirs = {b.name: b.vector.normalized() for b in bones if b.vector.length > 1e-6}
    hip = bones["Hips"].head.copy()
    torso_names = subtree(bones, "Spine1")
    head_z = bones["Head"].head.z
    torso_factor = (marks["head"][2] - hip.z) / (head_z - hip.z)
    print("FACTOR torso", round(torso_factor, 3))
    if not 0.4 <= torso_factor <= 2.5:
        raise SystemExit(f"torso factor {torso_factor}")
    scale_from(bones, torso_names, hip, torso_factor)
    for side, target in (("Left", marks["left_hand"]), ("Right", marks["right_hand"])):
        pivot = bones[f"{side}Arm"].head.copy()
        hand = bones[f"{side}Hand"].head
        print("HAND", side, "bone", round(hand.x, 3), "target", round(float(target[0]), 3), "pivot", round(pivot.x, 3))
        factor = (target[0] - pivot.x) / (hand.x - pivot.x)
        print("FACTOR", side, "arm", round(factor, 3))
        if not 0.2 <= factor <= 2.5:
            raise SystemExit(f"{side} arm factor {factor}")
        names = subtree(bones, f"{side}Arm")
        scale_from(bones, names, pivot, factor)
        hand = bones[f"{side}Hand"].head
        delta = Vector((0.0, target[1] - hand.y, target[2] - hand.z))
        for name in names:
            bones[name].head += delta
            bones[name].tail += delta
    for side, target in (("Left", marks["left_foot"]), ("Right", marks["right_foot"])):
        pivot = bones[f"{side}Leg"].head.copy()
        foot = bones[f"{side}Foot"].head
        factor = (target[2] - pivot.z) / (foot.z - pivot.z)
        print("FACTOR", side, "leg", round(factor, 3))
        if not 0.2 <= factor <= 2.5:
            raise SystemExit(f"{side} leg factor {factor}")
        names = subtree(bones, f"{side}Leg")
        scale_from(bones, names, pivot, factor)
        foot = bones[f"{side}Foot"].head
        delta = Vector((target[0] - foot.x, target[1] - foot.y, 0.0))
        for name in names:
            bones[name].head += delta
            bones[name].tail += delta
    worst = 0.0
    for name, direction in before_dirs.items():
        if name in ("LeftShoulder", "RightShoulder"):
            continue
        now = bones[name].vector.normalized()
        worst = max(worst, direction.angle(now, 0.0))
    bpy.ops.object.mode_set(mode="OBJECT")
    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 ("Hips", "Head", "LeftHand", "LeftFoot", "Chest"):
        p = heads[name]
        print(name, round(p.x, 3), round(p.y, 3), round(p.z, 3))
    print("AXIS_DRIFT_DEG", round(worst * 180 / pi, 3))
    if worst > 0.02:
        raise SystemExit("bone axes moved")
    mesh.parent = arm
    arm.name = "Armature"
    arm.data.name = "Armature"
    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)
    print("FIT", json.dumps({"bones": len(arm.data.bones), "fingers": fingers, "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")
    else:
        raise SystemExit("use template or fit")


if __name__ == "__main__":
    main()