formscout/analysis/laban.py

"""
Laban Movement Analysis — Effort factors from pose kinematics.

Computes the four Effort factors over the joints relevant to a screening test:

  - Space   (indirect 0 … direct 1)    — path directness of the leading joint
  - Weight  (light 0 … strong 1)       — motion-energy of the relevant joints
  - Time    (sustained 0 … sudden 1)   — impulsivity (peak vs mean speed)
  - Flow    (bound 0 … free 1)         — smoothness (inverse normalised jerk)

These are reproducible kinematic heuristics, not clinical LMA notation — the
report labels them as such. Distances are normalised by torso length so the
factors are scale-invariant across cameras. Pure function — no model, no I/O.
"""
from __future__ import annotations

import math

from formscout.agents.biomechanics import _get_joint
from formscout.analysis.relevant_joints import (
    COCO_NAMES, L_HIP, L_SHOULDER, R_HIP, R_SHOULDER, relevant_joints,
)

# Heuristic calibration references (movement is normalised to torso-lengths/sec).
_WEIGHT_REF = 1.0      # energy (bl/s)^2 giving ~0.63 weight
_TIME_LO, _TIME_HI = 1.5, 4.0   # peak/mean speed ratio mapped to [0, 1]
_FLOW_JERK_REF = 6.0   # normalised jerk (bl/s^3) giving ~0.37 flow

_LABELS = {
    "space": ("indirect", "direct"),
    "weight": ("light", "strong"),
    "time": ("sustained", "sudden"),
    "flow": ("bound", "free"),
}


def _torso_scale(frames) -> float:
    """Median shoulder-hip distance across frames; 1.0 if unmeasurable."""
    lengths = []
    for kps in frames:
        for sh, hip in ((L_SHOULDER, L_HIP), (R_SHOULDER, R_HIP)):
            a, b = _get_joint(kps, sh), _get_joint(kps, hip)
            if a and b:
                lengths.append(math.hypot(a[0] - b[0], a[1] - b[1]))
    if not lengths:
        return 1.0
    lengths.sort()
    med = lengths[len(lengths) // 2]
    return med if med > 1e-6 else 1.0


def _joint_kinematics(frames, joint_id: int, dt: float, scale: float) -> dict | None:
    """Speed/accel/jerk/directness for one joint trajectory (torso-length units)."""
    pts = [_get_joint(kps, joint_id) for kps in frames]
    valid = [(i, p) for i, p in enumerate(pts) if p is not None]
    if len(valid) < 3:
        return None

    speeds, path_len = [], 0.0
    for (i0, p0), (i1, p1) in zip(valid, valid[1:]):
        d = math.hypot(p1[0] - p0[0], p1[1] - p0[1]) / scale
        path_len += d
        gap = max(1, i1 - i0)
        speeds.append(d / (gap * dt))
    if not speeds:
        return None

    net = math.hypot(valid[-1][1][0] - valid[0][1][0],
                     valid[-1][1][1] - valid[0][1][1]) / scale
    directness = net / path_len if path_len > 1e-6 else 0.0

    accels = [abs(speeds[i + 1] - speeds[i]) / dt for i in range(len(speeds) - 1)]
    jerks = [abs(accels[i + 1] - accels[i]) / dt for i in range(len(accels) - 1)]

    mean_speed = sum(speeds) / len(speeds)
    peak_speed = max(speeds)
    return {
        "mean_speed": mean_speed,
        "peak_speed": peak_speed,
        "energy": sum(s * s for s in speeds) / len(speeds),
        "ratio": peak_speed / (mean_speed + 1e-6),
        "mean_jerk": (sum(jerks) / len(jerks)) if jerks else 0.0,
        "directness": min(1.0, max(0.0, directness)),
    }


def _clip01(x: float) -> float:
    return min(1.0, max(0.0, x))


def compute_laban(pose2d, test_name: str, fps: float) -> dict:
    """Return the four Effort factors, their labels, and body emphasis."""
    frames = pose2d.keypoints
    dt = 1.0 / fps if fps and fps > 0 else 1.0 / 30.0
    scale = _torso_scale(frames)
    joints = relevant_joints(test_name) or list(range(17))

    kin = {j: k for j in joints if (k := _joint_kinematics(frames, j, dt, scale))}
    if not kin:
        return {
            "effort": {"space": 0.0, "weight": 0.0, "time": 0.0, "flow": 0.0},
            "labels": {k: v[0] for k, v in _LABELS.items()},
            "body_emphasis": [],
            "notes": "insufficient motion to estimate Effort",
        }

    leader = max(kin, key=lambda j: kin[j]["mean_speed"])
    lead = kin[leader]

    weight = _clip01(1.0 - math.exp(-(sum(k["energy"] for k in kin.values()) / len(kin)) / _WEIGHT_REF))
    time = _clip01((lead["ratio"] - _TIME_LO) / (_TIME_HI - _TIME_LO))
    flow = _clip01(math.exp(-lead["mean_jerk"] / _FLOW_JERK_REF))
    space = _clip01(lead["directness"])

    effort = {"space": space, "weight": weight, "time": time, "flow": flow}
    labels = {k: _LABELS[k][1] if v >= 0.5 else _LABELS[k][0] for k, v in effort.items()}

    emphasis = sorted(kin.items(), key=lambda kv: kv[1]["mean_speed"], reverse=True)[:3]
    body_emphasis = [(COCO_NAMES.get(j, str(j)), round(k["mean_speed"], 3)) for j, k in emphasis]

    return {
        "effort": {k: round(v, 3) for k, v in effort.items()},
        "labels": labels,
        "body_emphasis": body_emphasis,
        "leading_joint": COCO_NAMES.get(leader, str(leader)),
        "notes": "kinematic Effort estimate (heuristic, not clinical LMA notation)",
    }