"""MIRPAMO stage 1: auto-rig a raw humanoid mesh.

Replicates Mixamo's auto-rigger locally:
  1. Landmark detection — instead of the user clicking colored dots on
     chin/wrists/elbows/knees/groin, we estimate those anchor points from the
     vertex cloud (T/A-pose assumed, character upright on the floor).
     A markers JSON can override any landmark.
  2. Skeletal template fitting — a mixamorig-named humanoid skeleton is
     stretched to the landmarks, so existing Mixamo clips retarget 1:1.
  3. Automatic skin weighting — Blender's bone-heat solver
     (parent with automatic weights), same spatial-falloff idea Mixamo uses.

Usage:
  blender -b --python bpy/autorig.py -- --in mesh.glb --out rigged.glb \
      [--markers markers.json] [--name Character]

markers.json (all optional, world-space [x,y,z] in Blender Z-up meters):
  {"chin": [...], "groin": [...], "l_wrist": [...], "r_wrist": [...],
   "l_elbow": [...], "r_elbow": [...], "l_knee": [...], "r_knee": [...]}
"""
import argparse
import json
import os
import sys

import bpy
from mathutils import Vector

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import (clean_scene, die, export_glb, find_meshes, import_any,
                    select_only, stage_args)


def gather_world_verts(meshes):
    """All vertices of all mesh objects, in world space."""
    depsgraph = bpy.context.evaluated_depsgraph_get()
    verts = []
    for obj in meshes:
        ev = obj.evaluated_get(depsgraph)
        mesh = ev.to_mesh()
        mw = ev.matrix_world
        verts.extend(mw @ v.co for v in mesh.vertices)
        ev.to_mesh_clear()
    if not verts:
        die("no vertices found in input mesh")
    return verts


def band(verts, zmin, zmax):
    return [v for v in verts if zmin <= v.z <= zmax]


def median(vals):
    s = sorted(vals)
    return s[len(s) // 2] if s else 0.0


def detect_landmarks(verts, overrides):
    """Estimate Mixamo-style anchor points from the vertex cloud.

    Assumes: upright character, +Z up, roughly symmetric about X=0,
    T- or A-pose. Heuristics follow standard human proportions, but arm
    heights are measured from the actual mesh so A-poses work too.
    """
    zs = [v.z for v in verts]
    z_min, z_max = min(zs), max(zs)
    H = z_max - z_min
    if H <= 0:
        die("mesh has zero height")
    cx = median([v.x for v in verts])
    cy = median([v.y for v in verts])

    L = {}

    def z_at(f):
        return z_min + f * H

    # --- torso column ---
    L["groin"] = Vector((cx, cy, z_at(0.52)))
    L["chin"] = Vector((cx, cy, z_at(0.87)))
    L["head_top"] = Vector((cx, cy, z_max))

    # --- arms: farthest-|x| vertex above the waist gives the hand tip;
    # wrist/elbow/shoulder sit at fixed fractions along that reach, and we
    # measure their actual heights from the mesh so drooping A-pose arms
    # land inside the geometry ---
    upper = band(verts, z_at(0.55), z_at(0.95))
    if not upper:
        upper = verts
    tip_l = max(upper, key=lambda v: v.x)
    tip_r = min(upper, key=lambda v: v.x)

    def arm_point(tip, frac):
        """Point at `frac` of the horizontal reach from center, at the
        median height/depth of mesh vertices near that x."""
        x = cx + (tip.x - cx) * frac
        tol = max(0.03 * H, 0.02)
        near = [v for v in upper if abs(v.x - x) < tol
                and (v.x - cx) * (tip.x - cx) >= 0]
        if len(near) < 3:
            return Vector((x, tip.y, tip.z))
        return Vector((x, median([v.y for v in near]),
                       median([v.z for v in near])))

    for side, tip in (("l", tip_l), ("r", tip_r)):
        L[f"{side}_hand_tip"] = Vector(tip)
        L[f"{side}_wrist"] = arm_point(tip, 0.82)
        L[f"{side}_elbow"] = arm_point(tip, 0.52)
        L[f"{side}_shoulder"] = arm_point(tip, 0.23)
        # keep the shoulder joint near the torso's height regardless of pose
        L[f"{side}_shoulder"].z = max(L[f"{side}_shoulder"].z, z_at(0.78))

    # --- legs: cluster low-body vertices left/right of center ---
    knee_band = band(verts, z_at(0.22), z_at(0.35))
    lknee = [v for v in knee_band if v.x > cx]
    rknee = [v for v in knee_band if v.x <= cx]
    leg_x_l = median([v.x for v in lknee]) if lknee else cx + 0.05 * H
    leg_x_r = median([v.x for v in rknee]) if rknee else cx - 0.05 * H
    leg_y_l = median([v.y for v in lknee]) if lknee else cy
    leg_y_r = median([v.y for v in rknee]) if rknee else cy

    L["l_knee"] = Vector((leg_x_l, leg_y_l, z_at(0.285)))
    L["r_knee"] = Vector((leg_x_r, leg_y_r, z_at(0.285)))
    L["l_ankle"] = Vector((leg_x_l, leg_y_l, z_at(0.05)))
    L["r_ankle"] = Vector((leg_x_r, leg_y_r, z_at(0.05)))

    # feet point toward the side of the mesh with more foot-level volume
    foot_band = band(verts, z_min, z_at(0.05))
    fy = median([v.y for v in foot_band]) if foot_band else cy
    toe_dir = -1.0 if fy < cy else 1.0
    foot_len = 0.12 * H
    for s, ax in (("l", leg_x_l), ("r", leg_x_r)):
        L[f"{s}_toe"] = Vector((ax, cy + toe_dir * foot_len, z_min + 0.01 * H))

    # --- apply user marker overrides (Mixamo's colored dots) ---
    for key, val in overrides.items():
        L[key] = Vector(val)

    L["_height"] = H
    L["_floor"] = z_min
    L["_center"] = Vector((cx, cy, 0))
    return L


def build_skeleton(L, name):
    """Build a mixamorig-named humanoid armature fitted to the landmarks."""
    arm_data = bpy.data.armatures.new(name + "_rig")
    arm_obj = bpy.data.objects.new(name + "_rig", arm_data)
    bpy.context.collection.objects.link(arm_obj)
    select_only([arm_obj], active=arm_obj)
    bpy.ops.object.mode_set(mode="EDIT")

    eb = arm_data.edit_bones
    P = "mixamorig:"

    groin, chin, top = L["groin"], L["chin"], L["head_top"]
    spine_bottom = groin
    neck_base = Vector((chin.x, chin.y, groin.z + (chin.z - groin.z) * 0.92))

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

    bones = {}

    def add(bname, head, tail, parent=None, connect=False):
        b = eb.new(P + bname)
        b.head, b.tail = head, tail
        if parent:
            b.parent = bones[parent]
            b.use_connect = connect
        bones[bname] = b
        return b

    # spine column
    s1 = lerp(spine_bottom, neck_base, 0.33)
    s2 = lerp(spine_bottom, neck_base, 0.66)
    add("Hips", groin, s1)
    add("Spine", s1, s2, "Hips", True)
    add("Spine1", s2, neck_base, "Spine", True)
    add("Spine2", neck_base, lerp(neck_base, chin, 0.6), "Spine1", True)
    add("Neck", lerp(neck_base, chin, 0.6), chin, "Spine2", True)
    add("Head", chin, top, "Neck", True)

    # arms
    for S, side in (("Left", "l"), ("Right", "r")):
        sh, el, wr, tip = (L[f"{side}_shoulder"], L[f"{side}_elbow"],
                           L[f"{side}_wrist"], L[f"{side}_hand_tip"])
        clav = lerp(Vector((groin.x, sh.y, sh.z)), sh, 0.35)
        add(f"{S}Shoulder", clav, sh, "Spine2")
        add(f"{S}Arm", sh, el, f"{S}Shoulder", True)
        add(f"{S}ForeArm", el, wr, f"{S}Arm", True)
        add(f"{S}Hand", wr, tip, f"{S}ForeArm", True)

    # legs
    for S, side in (("Left", "l"), ("Right", "r")):
        kn, an, toe = L[f"{side}_knee"], L[f"{side}_ankle"], L[f"{side}_toe"]
        hip = Vector((kn.x, groin.y, groin.z * 0.98 + L["_floor"] * 0.02))
        add(f"{S}UpLeg", hip, kn, "Hips")
        add(f"{S}Leg", kn, an, f"{S}UpLeg", True)
        add(f"{S}Foot", an, lerp(an, toe, 0.75), f"{S}Leg", True)
        add(f"{S}ToeBase", lerp(an, toe, 0.75), toe, f"{S}Foot", True)

    bpy.ops.object.mode_set(mode="OBJECT")
    return arm_obj


def point_to_segment_dist(p, a, b):
    ab = b - a
    denom = ab.length_squared
    if denom < 1e-12:
        return (p - a).length
    t = max(0.0, min(1.0, (p - a).dot(ab) / denom))
    return (p - (a + ab * t)).length


def rescue_unweighted(mesh, arm_obj):
    """Bone heat can leave disconnected islands (hands, toes, props) with
    zero weights — they'd stay frozen at bind pose. Assign such vertices to
    the nearest bone by spatial proximity (Mixamo's own fallback idea)."""
    bones = [(b.name,
              arm_obj.matrix_world @ b.head_local,
              arm_obj.matrix_world @ b.tail_local)
             for b in arm_obj.data.bones]
    vg_index = {vg.index: vg.name for vg in mesh.vertex_groups}
    mw = mesh.matrix_world
    rescued = 0
    for v in mesh.data.vertices:
        if any(g.weight > 1e-6 and g.group in vg_index for g in v.groups):
            continue
        p = mw @ v.co
        name = min(bones, key=lambda bn: point_to_segment_dist(p, bn[1], bn[2]))[0]
        vg = mesh.vertex_groups.get(name) or mesh.vertex_groups.new(name=name)
        vg.add([v.index], 1.0, "REPLACE")
        rescued += 1
    if rescued:
        print(f"[mirpamo] rescued {rescued} unweighted verts on {mesh.name} "
              "via nearest-bone proximity")


def skin(meshes, arm_obj):
    """Bone-heat automatic weights (Mixamo's spatial-falloff step)."""
    for m in meshes:
        # bone heat fails on unapplied transforms more often; apply them
        select_only([m], active=m)
        bpy.ops.object.transform_apply(location=True, rotation=True, scale=True)
    select_only(meshes + [arm_obj], active=arm_obj)
    bpy.ops.object.parent_set(type="ARMATURE_AUTO")
    for m in meshes:
        rescue_unweighted(m, arm_obj)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--in", dest="inp", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--markers", default=None)
    ap.add_argument("--name", default=None)
    args = ap.parse_args(stage_args())

    overrides = {}
    if args.markers:
        with open(args.markers) as f:
            overrides = json.load(f)

    name = args.name or os.path.splitext(os.path.basename(args.inp))[0]

    clean_scene()
    objs = import_any(args.inp)
    meshes = find_meshes(objs)
    if not meshes:
        die("input contains no mesh")

    verts = gather_world_verts(meshes)
    L = detect_landmarks(verts, overrides)
    print(f"[mirpamo] mesh height {L['_height']:.3f}m, "
          f"{len(verts)} verts, landmarks detected")

    arm_obj = build_skeleton(L, name)
    skin(meshes, arm_obj)

    # sanity: every mesh must now have an armature modifier
    for m in meshes:
        if not any(mod.type == "ARMATURE" for mod in m.modifiers):
            die(f"automatic weighting failed for {m.name} "
                "(bone heat could not solve — mesh may need cleanup)")

    export_glb(args.out, objects=[arm_obj] + meshes)
    print("[mirpamo] autorig OK")


main()