#!/usr/bin/env python
# dance_profile.py — Blender-headless dance motion profiler.
#
# Samples world-space pose-bone positions every frame of a dance clip's GLB, then
# derives a per-beat motion-energy timeline + accent / phrase / totals metrics. The
# method generalises the one-off analysis behind docs/dances/dance1.md (which sampled
# world-space bone positions per frame) into a reusable script that ALSO emits a beat
# timeline at a given musical BPM.
#
# Usage (Blender bundles Python + numpy — no pip installs):
#   /Applications/Blender.app/Contents/MacOS/Blender --background \
#       --python tools/dance_profile.py -- --glb <path.glb> --bpm 100 --out <profile.json>
#
# Output JSON schema: see docs/dances/dance1.md / task spec. Units are METERS (glTF is
# always meters; the retargeted Quaternius rigs bake object transform so Blender world
# == data space, Z-up). Pair groups (hands, forearms, ...) report the MEAN of the two
# sides' path lengths — this reproduces dance1.md ("hands ~49-50 m each") and the
# schema example (hands_m: 49.5); the pair sum would be ~2x.

import bpy
import sys
import os
import json
import numpy as np


# ---------------------------------------------------------------------------
# Bone name resolution
#
# Rigs in this repo are Quaternius UE-named (pelvis / Head / hand_l / lowerarm_l /
# calf_l ...). We also tolerate CC_Base_* (CC_Base_L_Hand / CC_Base_L_Calf ...) and
# Mixamo B-* (B-hand.L / B-shin.L ...) so the same script works on pre-retarget
# source rigs. Matching is case-insensitive "contains", first candidate wins.
# ---------------------------------------------------------------------------

# Center-line bones (no L/R).
CENTER_CANDIDATES = {
    "pelvis": ["pelvis", "hips", "hip"],
    "head":   ["head"],
}

# Paired bones: group -> (left_candidates, right_candidates).
SIDE_CANDIDATES = {
    "hand":     (["hand_l", "hand.l", "l_hand", "cc_base_l_hand", "b-hand.l", "hand_left"],
                 ["hand_r", "hand.r", "r_hand", "cc_base_r_hand", "b-hand.r", "hand_right"]),
    "forearm":  (["lowerarm_l", "lowerarm.l", "forearm_l", "forearm.l", "l_forearm",
                  "cc_base_l_forearm", "b-forearm.l", "lowerarm_left"],
                 ["lowerarm_r", "lowerarm.r", "forearm_r", "forearm.r", "r_forearm",
                  "cc_base_r_forearm", "b-forearm.r", "lowerarm_right"]),
    "upperarm": (["upperarm_l", "upperarm.l", "l_upperarm", "cc_base_l_upperarm",
                  "b-upperarm.l", "upperarm_left"],
                 ["upperarm_r", "upperarm.r", "r_upperarm", "cc_base_r_upperarm",
                  "b-upperarm.r", "upperarm_right"]),
    "foot":     (["foot_l", "foot.l", "l_foot", "cc_base_l_foot", "b-foot.l", "foot_left"],
                 ["foot_r", "foot.r", "r_foot", "cc_base_r_foot", "b-foot.r", "foot_right"]),
    "shin":     (["calf_l", "calf.l", "shin_l", "shin.l", "l_calf", "cc_base_l_calf",
                  "b-shin.l", "b-calf.l", "calf_left", "shin_left"],
                 ["calf_r", "calf.r", "shin_r", "shin.r", "r_calf", "cc_base_r_calf",
                  "b-shin.r", "b-calf.r", "calf_right", "shin_right"]),
}

# All logical group keys the metrics expect.
ALL_GROUPS = ["pelvis", "head",
              "hand_l", "hand_r", "forearm_l", "forearm_r",
              "upperarm_l", "upperarm_r", "foot_l", "foot_r",
              "shin_l", "shin_r"]


def resolve_by_candidates(pose_bone_names, candidates):
    """Return the first pose-bone name whose lowercased form contains any candidate
    (candidates tried in priority order). None if nothing matches."""
    lowered = [(n, n.lower()) for n in pose_bone_names]
    for cand in candidates:
        for name, low in lowered:
            if cand in low:
                return name
    return None


def resolve_bones(armature, warnings):
    """Map every logical group to a pose-bone name (or None). Appends a warning per
    missing group."""
    pose_names = [pb.name for pb in armature.pose.bones]
    resolved = {}

    for group, cands in CENTER_CANDIDATES.items():
        resolved[group] = resolve_by_candidates(pose_names, cands)

    for group, (left_c, right_c) in SIDE_CANDIDATES.items():
        resolved[f"{group}_l"] = resolve_by_candidates(pose_names, left_c)
        resolved[f"{group}_r"] = resolve_by_candidates(pose_names, right_c)

    for g in ALL_GROUPS:
        if resolved.get(g) is None:
            warnings.append(f"bone group not found: {g}")
    return resolved


# ---------------------------------------------------------------------------
# Scene / action setup
# ---------------------------------------------------------------------------

def parse_argv():
    """Args after the `--` separator that Blender leaves for the script."""
    argv = sys.argv[sys.argv.index("--") + 1:] if "--" in sys.argv else []
    args = {}
    i = 0
    while i < len(argv):
        key = argv[i]
        if key.startswith("--"):
            args[key] = argv[i + 1] if i + 1 < len(argv) else ""
            i += 2
        else:
            i += 1
    return args


def find_armature():
    arms = [o for o in bpy.data.objects if o.type == "ARMATURE"]
    if not arms:
        raise RuntimeError("no armature object in imported GLB")
    if len(arms) > 1:
        raise RuntimeError(f"expected exactly one armature, found {len(arms)}: "
                           f"{[a.name for a in arms]}")
    return arms[0]


def get_action(armature):
    """Active action on the armature, else the first action in bpy.data.actions."""
    ad = armature.animation_data
    if ad and ad.action:
        return ad.action
    if bpy.data.actions:
        a = bpy.data.actions[0]
        armature.animation_data_create()
        armature.animation_data.action = a
        return a
    raise RuntimeError("armature has no animation data and bpy.data.actions is empty")


# ---------------------------------------------------------------------------
# Sampling
# ---------------------------------------------------------------------------

def sample_positions(armature, resolved, frame_list):
    """World-space translation of every resolved bone across frame_list.
    Returns {group: np.ndarray (N,3)} for resolved groups only."""
    scene = bpy.context.scene
    mw = armature.matrix_world
    groups = {g: name for g, name in resolved.items() if name is not None}
    out = {g: np.empty((len(frame_list), 3), dtype=np.float64) for g in groups}

    for fi, f in enumerate(frame_list):
        scene.frame_set(f)
        for g, name in groups.items():
            pb = armature.pose.bones[name]
            t = (mw @ pb.matrix).translation
            out[g][fi, 0] = t.x
            out[g][fi, 1] = t.y
            out[g][fi, 2] = t.z
    return out


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

def path_length(pos):
    d = np.diff(pos, axis=0)
    return float(np.sum(np.linalg.norm(d, axis=1)))


def per_interval_speed(pos, fps):
    """Speed (m/s) for each inter-frame interval. Length N-1."""
    d = np.diff(pos, axis=0)
    return np.linalg.norm(d, axis=1) * fps


def local_min_indices(sig):
    return [i for i in range(1, len(sig) - 1)
            if sig[i] < sig[i - 1] and sig[i] < sig[i + 1]]


def local_max_indices(sig):
    return [i for i in range(1, len(sig) - 1)
            if sig[i] > sig[i - 1] and sig[i] > sig[i + 1]]


def pearson(a, b):
    a = a - a.mean()
    b = b - b.mean()
    da = np.sqrt(np.dot(a, a))
    db = np.sqrt(np.dot(b, b))
    if da < 1e-12 or db < 1e-12:
        return 0.0
    return float(np.dot(a, b) / (da * db))


def compute_profile(name, glb_arg, positions, resolved, fps, frames, f0, f1, bpm, warnings):
    N = frames  # number of sampled frames
    duration_s = frames / fps
    beats = int(round(duration_s * bpm / 60.0))
    if beats < 1:
        beats = 1
    fpb = frames / float(beats)  # frames per beat (float)

    dt_index = lambda sample_idx: int(min(beats - 1, max(0, sample_idx / fpb)))

    # ---- energy_per_beat: mean world speed over hands+feet+pelvis --------------
    energy_groups = ["hand_l", "hand_r", "foot_l", "foot_r", "pelvis"]
    energy_groups = [g for g in energy_groups if g in positions]
    if not energy_groups:
        warnings.append("no energy bones (hands/feet/pelvis) resolved — energy timeline zero")

    energy_per_beat = []
    for b in range(beats):
        i_start = int(np.floor(b * fpb))
        i_end = int(np.floor((b + 1) * fpb))
        i_start = max(0, min(i_start, N - 1))
        i_end = max(i_start + 1, min(i_end, N - 1))  # need >=1 delta
        window_s = (i_end - i_start) / fps
        bone_speeds = []
        for g in energy_groups:
            seg = positions[g][i_start:i_end + 1]
            dist = path_length(seg)
            bone_speeds.append(dist / window_s if window_s > 0 else 0.0)
        energy_per_beat.append(float(np.mean(bone_speeds)) if bone_speeds else 0.0)

    emin, emax = float(min(energy_per_beat)), float(max(energy_per_beat))
    if emax - emin > 1e-12:
        energy_norm = [(e - emin) / (emax - emin) for e in energy_per_beat]
    else:
        energy_norm = [0.0 for _ in energy_per_beat]

    # ---- accent beats -------------------------------------------------------
    pelvis_drop_beats = []
    if "pelvis" in positions:
        pz = positions["pelvis"][:, 2]
        mean_z, std_z = float(pz.mean()), float(pz.std())
        thresh_min = mean_z - 0.5 * std_z
        for i in local_min_indices(pz):
            if pz[i] < thresh_min:
                pelvis_drop_beats.append(dt_index(i))
    else:
        warnings.append("pelvis missing — pelvis-drop accents skipped")
    pelvis_drop_beats = sorted(set(pelvis_drop_beats))

    hand_peak_beats = []
    if "hand_l" in positions or "hand_r" in positions:
        hs = np.zeros(max(0, N - 1), dtype=np.float64)
        for g in ("hand_l", "hand_r"):
            if g in positions:
                hs = hs + per_interval_speed(positions[g], fps)
        if len(hs) > 2:
            mean_hs, std_hs = float(hs.mean()), float(hs.std())
            thresh_max = mean_hs + 1.0 * std_hs
            for i in local_max_indices(hs):
                if hs[i] > thresh_max:
                    hand_peak_beats.append(dt_index(i))
    else:
        warnings.append("both hands missing — hand-peak accents skipped")
    hand_peak_beats = sorted(set(hand_peak_beats))

    accent_beats = sorted(set(pelvis_drop_beats) | set(hand_peak_beats))

    # ---- phrase_beats: dominant period via autocorrelation of energy_norm ----
    en = np.array(energy_norm, dtype=np.float64)
    best_lag, best_r = None, -1.0
    for k in range(2, max(2, beats // 2) + 1):
        if k >= len(en):
            break
        r = pearson(en[:-k], en[k:])
        if r > best_r:
            best_r, best_lag = r, k
    phrase_beats = int(best_lag) if (best_lag is not None and best_r >= 0.3) else None

    # ---- totals -------------------------------------------------------------
    def pair_path_mean(g):
        ls = path_length(positions[g + "_l"]) if (g + "_l") in positions else None
        rs = path_length(positions[g + "_r"]) if (g + "_r") in positions else None
        vals = [v for v in (ls, rs) if v is not None]
        return float(np.mean(vals)) if vals else None

    def lr_symmetry(g):
        ls = path_length(positions[g + "_l"]) if (g + "_l") in positions else 0.0
        rs = path_length(positions[g + "_r"]) if (g + "_r") in positions else 0.0
        if max(ls, rs) < 1e-9:
            return 1.0
        return float(min(ls, rs) / max(ls, rs))

    totals = {
        "hands_m":     pair_path_mean("hand"),
        "forearms_m":  pair_path_mean("forearm"),
        "upper_arms_m": pair_path_mean("upperarm"),
        "feet_m":      pair_path_mean("foot"),
        "shins_m":     pair_path_mean("shin"),
        "pelvis_m":    path_length(positions["pelvis"]) if "pelvis" in positions else None,
        "head_m":      path_length(positions["head"]) if "head" in positions else None,
        "lr_symmetry_hands": lr_symmetry("hand"),
        "lr_symmetry_feet":  lr_symmetry("foot"),
    }
    if "pelvis" in positions:
        p = positions["pelvis"]
        # Blender is Z-up: vertical range = Z; the horizontal footprint = X,Y.
        # (The spec's "[x_range, z_range]" is glTF Y-up notation — glTF's Z maps to
        # Blender Y on import — so this reproduces the schema's [1.2, 1.7] footing.)
        totals["pelvis_vertical_range_m"] = float(p[:, 2].max() - p[:, 2].min())
        totals["footprint_m"] = [float(p[:, 0].max() - p[:, 0].min()),
                                 float(p[:, 1].max() - p[:, 1].min())]
    else:
        totals["pelvis_vertical_range_m"] = None
        totals["footprint_m"] = [None, None]

    profile = {
        "name": name,
        "glb": glb_arg,
        "duration_s": round(duration_s, 4),
        "fps": fps,
        "frames": frames,
        "bpm": bpm,
        "beats": beats,
        "energy_per_beat": [round(e, 4) for e in energy_per_beat],
        "energy_norm": [round(e, 4) for e in energy_norm],
        "accent_beats": accent_beats,
        "pelvis_drop_beats": pelvis_drop_beats,
        "hand_peak_beats": hand_peak_beats,
        "phrase_beats": phrase_beats,
        "totals": totals,
        "warnings": warnings,
    }
    return profile


def print_summary(profile):
    accent = profile["accent_beats"]
    top = accent[:8]
    print("[profile] %s : %.2fs, %d frames @ %d fps"
          % (profile["name"], profile["duration_s"], profile["frames"], profile["fps"]))
    print("[profile] %d beats @ %d bpm (energy mean %.2f, peak %.2f)"
          % (profile["beats"], profile["bpm"],
             float(np.mean(profile["energy_per_beat"])),
             float(max(profile["energy_per_beat"]))))
    print("[profile] accent beats (%d): %s%s"
          % (len(accent), top, " ..." if len(accent) > 8 else ""))
    print("[profile] pelvis-drops: %s | hand-peaks: %s | phrase: %s beats"
          % (profile["pelvis_drop_beats"], profile["hand_peak_beats"],
             profile["phrase_beats"] if profile["phrase_beats"] is not None else "null"))
    w = profile["warnings"]
    print("[profile] warnings (%d): %s" % (len(w), w if w else "none"))


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main():
    args = parse_argv()
    glb_arg = args.get("--glb")
    bpm = int(args.get("--bpm", "100"))
    out = args.get("--out")
    if not glb_arg or not out:
        raise SystemExit("usage: dance_profile.py -- --glb <path> --bpm <n> --out <path>")

    bpy.ops.wm.read_factory_settings(use_empty=True)
    bpy.ops.import_scene.gltf(filepath=glb_arg)
    bpy.context.view_layer.update()

    armature = find_armature()
    action = get_action(armature)
    f0, f1 = action.frame_range
    f0, f1 = int(f0), int(f1)
    # Scene fps from the imported scene (glTF carries no per-action fps metadata;
    # the importer bakes the clip at the scene rate, so frames/fps == source duration).
    fps = bpy.context.scene.render.fps

    warnings = []
    resolved = resolve_bones(armature, warnings)
    frame_list = list(range(f0, f1 + 1))
    frames = len(frame_list)

    positions = sample_positions(armature, resolved, frame_list)

    name = os.path.splitext(os.path.basename(glb_arg))[0]
    profile = compute_profile(name, glb_arg, positions, resolved, fps, frames,
                              f0, f1, bpm, warnings)

    os.makedirs(os.path.dirname(os.path.abspath(out)), exist_ok=True)
    with open(out, "w") as fh:
        json.dump(profile, fh, indent=2)

    # Sanity line so we can confirm units are meters on the first run.
    if "pelvis" in positions:
        pz = positions["pelvis"][:, 2]
        print("[profile] pelvis z: mean %.3f m, range %.3f m (units check)"
              % (float(pz.mean()), float(pz.max() - pz.min())))
    print_summary(profile)
    print("[profile] WROTE %s" % out)


if __name__ == "__main__":
    main()