"""Joint angle computation from pose landmarks.

Supports multiple computation methods via a registry:
    - "sagittal_vertical_axis" (default): vertical-axis reference method.
      Uses atan2 of the thigh vector relative to the trunk vector,
      providing natural 0° in standing posture and avoiding ±180°
      discontinuities common in 3-point methods.
      Ref: Davis RB, Ounpuu S, Tyburski D, Gage JR. A gait analysis
      data collection and reduction technique. Hum Mov Sci.
      1991;10(5):575-587. doi:10.1016/0167-9457(91)90046-Z

    - "sagittal_classic": classical 3-point interior angle method.
      Ref: Kadaba MP, Ramakrishnan HK, Wootten ME. Measurement of
      lower extremity kinematics during level walking. J Orthop Res.
      1990;8(3):383-392. doi:10.1002/jor.1100080310

Each method is a function(frame, model) -> dict of angles for one frame.
New methods can be registered via register_angle_method().

Angles convention (ISB recommendations):
    - Hip: positive = flexion, negative = extension
    - Knee: 0 = full extension, positive = flexion
    - Ankle: positive = dorsiflexion, negative = plantarflexion
    - Trunk: positive = forward lean
    - Pelvis tilt: positive = right side up

References for conventions:
    Wu G, Siegler S, Allard P, et al. ISB recommendation on definitions
    of joint coordinate system of various joints for the reporting of
    human joint motion — part I: ankle, hip, and spine. J Biomech.
    2002;35(4):543-548. doi:10.1016/S0021-9290(01)00222-6

Correction factor:
    The default correction_factor=0.8 compensates for 2D projection
    overestimation of sagittal-plane joint ROM.
    Ref: Mündermann L, Corazza S, Andriacchi TP. The evolution of
    methods for the capture of human movement leading to markerless
    motion capture for biomechanical applications. J Neuroeng Rehabil.
    2006;3:6. doi:10.1186/1743-0003-3-6
"""

import copy
import logging
from typing import Callable, Dict, Optional

import numpy as np


logger = logging.getLogger(__name__)


# ── Geometry helpers ─────────────────────────────────────────────────


def _get_xy(frame: dict, name: str) -> Optional[np.ndarray]:
    """Extract [x, y] array from a frame's landmarks dict."""
    lm = frame.get("landmarks", {}).get(name)
    if lm is None:
        return None
    x, y = lm.get("x"), lm.get("y")
    if x is None or y is None:
        return None
    try:
        if np.isnan(float(x)) or np.isnan(float(y)):
            return None
    except (TypeError, ValueError):
        pass
    return np.array([x, y])


def _angle_between(v1: np.ndarray, v2: np.ndarray) -> float:
    """Angle between two vectors in degrees [0, 180]."""
    n1 = np.linalg.norm(v1)
    n2 = np.linalg.norm(v2)
    if n1 < 1e-10 or n2 < 1e-10:
        return np.nan
    cos_a = np.clip(np.dot(v1, v2) / (n1 * n2), -1.0, 1.0)
    return np.degrees(np.arccos(cos_a))


def _trunk_angle(left_shoulder: np.ndarray, right_shoulder: np.ndarray,
                 left_hip: np.ndarray, right_hip: np.ndarray) -> float:
    """Trunk lean relative to vertical. Positive = forward lean."""
    shoulder_center = (left_shoulder + right_shoulder) / 2
    hip_center = (left_hip + right_hip) / 2
    trunk = shoulder_center - hip_center
    vertical = np.array([0.0, -1.0])  # up in image coords
    angle = _angle_between(vertical, trunk)
    return angle if trunk[0] > 0 else -angle


def _pelvis_tilt(left_hip: np.ndarray, right_hip: np.ndarray) -> float:
    """Pelvis tilt. Positive = right side up.

    Returns NaN in lateral view (hips overlap < 2%).
    """
    pelvis = right_hip - left_hip
    dist = np.linalg.norm(pelvis)
    if dist < 0.02:
        return np.nan
    horizontal = np.array([1.0, 0.0])
    angle = _angle_between(horizontal, pelvis)
    return angle if pelvis[1] < 0 else -angle


def _has_detected_foot(lm: dict, side: str) -> bool:
    """Return True if the heel for *side* is already a real detected landmark.

    A detected heel will have visibility > 0.3.  When the enrichment step
    in extract.py has injected real foot landmarks from Sapiens or RTMW
    auxiliary data, those landmarks carry the model's confidence score
    (typically > 0.8).  Geometric estimates are assigned visibility = 0.5
    but should NOT prevent re-estimation when they come from a prior
    pipeline run; the check here is intentionally > 0.3 so that both
    detected and estimated landmarks are accepted, but missing/NaN ones
    are not.
    """
    heel = lm.get(f"{side}_HEEL")
    fi = lm.get(f"{side}_FOOT_INDEX")
    if heel is None or fi is None:
        return False
    hx = heel.get("x")
    fx = fi.get("x")
    hv = heel.get("visibility", 0.0)
    fv = fi.get("visibility", 0.0)
    if hx is None or fx is None:
        return False
    if isinstance(hx, float) and np.isnan(hx):
        return False
    if isinstance(fx, float) and np.isnan(fx):
        return False
    return hv > 0.3 and fv > 0.3


def _estimate_foot_landmarks(frame: dict) -> dict:
    """Estimate HEEL and FOOT_INDEX from ANKLE for COCO models.

    If real detected foot landmarks are already present (injected by
    ``_enrich_foot_landmarks`` in extract.py from Sapiens / RTMW
    auxiliary data), those landmarks are kept and estimation is skipped
    for that side.
    """
    lm = frame.get("landmarks", {})

    # Check if any side actually needs estimation before deep-copying
    needs_estimation = False
    for side in ("LEFT", "RIGHT"):
        if _has_detected_foot(lm, side):
            continue
        ankle = _get_xy(frame, f"{side}_ANKLE")
        knee = _get_xy(frame, f"{side}_KNEE")
        if ankle is None or knee is None:
            continue
        needs_estimation = True
        break

    if not needs_estimation:
        return frame

    frame = copy.deepcopy(frame)
    lm = frame.get("landmarks", {})

    for side in ("LEFT", "RIGHT"):
        # Skip estimation when real detected foot landmarks exist
        if _has_detected_foot(lm, side):
            continue

        ankle = _get_xy(frame, f"{side}_ANKLE")
        knee = _get_xy(frame, f"{side}_KNEE")
        if ankle is None or knee is None:
            continue

        dx = ankle[0] - knee[0]
        dy = ankle[1] - knee[1]
        length = np.sqrt(dx ** 2 + dy ** 2)
        if length == 0:
            continue
        foot_len = length * 0.25

        heel_name = f"{side}_HEEL"
        heel = lm.get(heel_name)
        heel_is_nan = False
        try:
            heel_is_nan = heel is not None and heel.get("x") is not None and np.isnan(float(heel["x"]))
        except (TypeError, ValueError):
            pass
        if heel is None or heel.get("x") is None or heel_is_nan:
            lm[heel_name] = {
                "x": float(ankle[0] + dx / length * foot_len * 0.3),
                "y": float(ankle[1] + dy / length * foot_len * 0.3 + foot_len * 0.1),
                "visibility": 0.5,
            }

        fi_name = f"{side}_FOOT_INDEX"
        fi = lm.get(fi_name)
        fi_is_nan = False
        try:
            fi_is_nan = fi is not None and fi.get("x") is not None and np.isnan(float(fi["x"]))
        except (TypeError, ValueError):
            pass
        if fi is None or fi.get("x") is None or fi_is_nan:
            lm[fi_name] = {
                "x": float(ankle[0] - dx / length * foot_len * 0.5),
                "y": float(ankle[1] + foot_len * 0.15),
                "visibility": 0.5,
            }

    frame["landmarks"] = lm
    return frame


def _get_foot_index_from_toes(frame: dict, side: str) -> Optional[np.ndarray]:
    """Compute foot index as midpoint of big_toe and small_toe when available.

    Falls back to the existing FOOT_INDEX landmark if toe data is not
    present.  This gives a more accurate foot reference point when
    real toe landmarks have been injected from Sapiens or RTMW
    auxiliary data.
    """
    big_toe = _get_xy(frame, f"{side}_BIG_TOE")
    small_toe = _get_xy(frame, f"{side}_SMALL_TOE")
    if big_toe is not None and small_toe is not None:
        return (big_toe + small_toe) / 2.0
    # Fallback to standard FOOT_INDEX
    return _get_xy(frame, f"{side}_FOOT_INDEX")


def _unwrap_angles(values: list) -> list:
    """Unwrap angle series to remove ±180 discontinuities.

    Does NOT recenter on median -- the absolute anatomical reference
    must be preserved.  Recentering, if needed, is handled solely by
    neutral calibration (which by default only touches ankle joints).
    """
    arr = np.array(values, dtype=float)
    mask = ~np.isnan(arr)
    if mask.sum() < 2:
        return values
    unwrapped = np.degrees(np.unwrap(np.radians(arr[mask])))
    result = arr.copy()
    result[mask] = unwrapped
    return result.tolist()


def _detect_walking_direction(data: dict) -> str:
    """Detect walking direction from hip center horizontal displacement.

    Analyses the mean horizontal displacement of the hip center
    (average of LEFT_HIP and RIGHT_HIP x-coordinates) across all
    frames.  If x increases over time the subject walks left-to-right
    (in image coordinates); if it decreases the subject walks
    right-to-left.

    Parameters
    ----------
    data : dict
        Pivot JSON dict with ``frames`` populated.

    Returns
    -------
    str
        ``"left_to_right"`` or ``"right_to_left"``.
    """
    frames = data.get("frames", [])
    if len(frames) < 2:
        return "left_to_right"

    xs = []
    for frame in frames:
        if frame.get("confidence", 1.0) < 0.1:
            continue
        lm = frame.get("landmarks", {})
        lh = lm.get("LEFT_HIP")
        rh = lm.get("RIGHT_HIP")
        if (lh is not None and rh is not None
                and lh.get("x") is not None and rh.get("x") is not None
                and not np.isnan(lh["x"]) and not np.isnan(rh["x"])):
            xs.append((lh["x"] + rh["x"]) / 2.0)

    if len(xs) < 2:
        return "left_to_right"

    # Use the overall displacement (last quarter vs first quarter) to be
    # robust to noise.
    n = len(xs)
    q = max(1, n // 4)
    mean_start = float(np.mean(xs[:q]))
    mean_end = float(np.mean(xs[-q:]))

    return "left_to_right" if mean_end >= mean_start else "right_to_left"


def _extract_landmark_positions(frame: dict) -> dict:
    """Extract landmark positions for gaitkit compatibility."""
    lp = {}
    for name in ["left_hip", "right_hip", "left_knee", "right_knee",
                  "left_ankle", "right_ankle"]:
        mp_name = name.upper()
        pt = _get_xy(frame, mp_name)
        vis = frame.get("landmarks", {}).get(mp_name, {}).get("visibility", 0.0)
        if pt is not None:
            lp[name] = [float(pt[0]), float(pt[1]), float(vis)]
    return lp


# ── Method: sagittal_vertical_axis (default) ────────────────────────


def _method_sagittal_vertical_axis(frame: dict, model: str) -> dict:
    """Vertical-axis reference method (gaitopensim).

    Hip angle = atan2(thigh_from_vertical) - atan2(trunk_from_vertical).
    Natural 0° in standing posture, no discontinuities.
    """
    needs_foot = model != "mediapipe"
    f = _estimate_foot_landmarks(frame) if needs_foot else frame
    result = {"frame_idx": frame["frame_idx"]}

    l_shoulder = _get_xy(f, "LEFT_SHOULDER")
    r_shoulder = _get_xy(f, "RIGHT_SHOULDER")
    l_hip = _get_xy(f, "LEFT_HIP")
    r_hip = _get_xy(f, "RIGHT_HIP")

    if l_shoulder is not None and r_shoulder is not None and l_hip is not None and r_hip is not None:
        hip_center = (l_hip + r_hip) / 2
        shoulder_center = (l_shoulder + r_shoulder) / 2
        trunk_vec = hip_center - shoulder_center
        result["trunk_angle"] = _trunk_angle(l_shoulder, r_shoulder, l_hip, r_hip)
        result["pelvis_tilt"] = _pelvis_tilt(l_hip, r_hip)
    else:
        trunk_vec = None
        result["trunk_angle"] = np.nan
        result["pelvis_tilt"] = np.nan

    result["landmark_positions"] = _extract_landmark_positions(f)

    for side, prefix in [("L", "LEFT"), ("R", "RIGHT")]:
        hip = _get_xy(f, f"{prefix}_HIP")
        knee = _get_xy(f, f"{prefix}_KNEE")
        ankle = _get_xy(f, f"{prefix}_ANKLE")
        # Use real toe midpoint when available, else standard FOOT_INDEX
        foot = _get_foot_index_from_toes(f, prefix)

        # Hip: vertical axis method
        # Positive = flexion (thigh forward of trunk), negative = extension
        if hip is not None and knee is not None and trunk_vec is not None:
            thigh_vector = knee - hip
            thigh_angle = np.arctan2(thigh_vector[0], thigh_vector[1])
            trunk_angle = np.arctan2(trunk_vec[0], trunk_vec[1])
            raw = np.degrees(trunk_angle - thigh_angle)
            # Normalize to [-180, 180] to avoid wrap-around artifacts
            result[f"hip_{side}"] = float(((raw + 180) % 360) - 180)
        else:
            result[f"hip_{side}"] = np.nan

        # Knee: signed flexion angle (positive = flexion, negative = hyperextension)
        if hip is not None and knee is not None and ankle is not None:
            thigh = hip - knee
            shank = ankle - knee
            unsigned = _angle_between(thigh, shank)
            # Cross product sign: positive when shank is "behind" thigh (flexion)
            cross = float(thigh[0] * shank[1] - thigh[1] * shank[0])
            result[f"knee_{side}"] = (180.0 - unsigned) if cross >= 0 else -(180.0 - unsigned)
        else:
            result[f"knee_{side}"] = np.nan

        # Ankle: Formula E — both vectors from ANKLE (the anatomical pivot).
        # shank = KNEE - ANKLE, foot_seg = FOOT_INDEX - ANKLE.
        # Angle = 90° minus the unsigned angle between the two vectors,
        # with sign determined by the 2D cross product.
        # Positive = dorsiflexion, negative = plantarflexion.
        if knee is not None and ankle is not None and foot is not None:
            shank = knee - ankle
            foot_seg = foot - ankle
            denom = np.linalg.norm(shank) * np.linalg.norm(foot_seg) + 1e-12
            cos_val = np.clip(np.dot(shank, foot_seg) / denom, -1, 1)
            unsigned = np.degrees(np.arccos(cos_val))
            cross = float(shank[0] * foot_seg[1] - shank[1] * foot_seg[0])
            result[f"ankle_{side}"] = (90.0 - unsigned) if cross >= 0 else -(90.0 - unsigned)
        else:
            result[f"ankle_{side}"] = np.nan

    return result


# ── Method: sagittal_classic ─────────────────────────────────────────


def _method_sagittal_classic(frame: dict, model: str) -> dict:
    """Classical 3-point interior angle method.

    Hip = 180 - angle(shoulder-hip-knee). 0° = aligned, positive = flexion.
    Knee = 180 - angle(hip-knee-ankle). 0° = full extension.
    Ankle = 90 - angle(knee-ankle-foot). dorsi/plantarflexion.

    Simpler than vertical-axis but has a non-zero offset in standing.
    """
    needs_foot = model != "mediapipe"
    f = _estimate_foot_landmarks(frame) if needs_foot else frame
    result = {"frame_idx": frame["frame_idx"]}

    l_shoulder = _get_xy(f, "LEFT_SHOULDER")
    r_shoulder = _get_xy(f, "RIGHT_SHOULDER")
    l_hip = _get_xy(f, "LEFT_HIP")
    r_hip = _get_xy(f, "RIGHT_HIP")

    if l_shoulder is not None and r_shoulder is not None and l_hip is not None and r_hip is not None:
        result["trunk_angle"] = _trunk_angle(l_shoulder, r_shoulder, l_hip, r_hip)
        result["pelvis_tilt"] = _pelvis_tilt(l_hip, r_hip)
    else:
        result["trunk_angle"] = np.nan
        result["pelvis_tilt"] = np.nan

    result["landmark_positions"] = _extract_landmark_positions(f)

    for side, prefix in [("L", "LEFT"), ("R", "RIGHT")]:
        shoulder = _get_xy(f, f"{prefix}_SHOULDER")
        hip = _get_xy(f, f"{prefix}_HIP")
        knee = _get_xy(f, f"{prefix}_KNEE")
        ankle = _get_xy(f, f"{prefix}_ANKLE")
        # Use real toe midpoint when available, else standard FOOT_INDEX
        foot = _get_foot_index_from_toes(f, prefix)

        # Hip: 180 - angle(shoulder, hip, knee)
        if shoulder is not None and hip is not None and knee is not None:
            v1 = shoulder - hip
            v2 = knee - hip
            result[f"hip_{side}"] = 180.0 - _angle_between(v1, v2)
        else:
            result[f"hip_{side}"] = np.nan

        # Knee: signed flexion angle (positive = flexion, negative = hyperextension)
        if hip is not None and knee is not None and ankle is not None:
            thigh = hip - knee
            shank = ankle - knee
            unsigned = _angle_between(thigh, shank)
            cross = float(thigh[0] * shank[1] - thigh[1] * shank[0])
            result[f"knee_{side}"] = (180.0 - unsigned) if cross >= 0 else -(180.0 - unsigned)
        else:
            result[f"knee_{side}"] = np.nan

        # Ankle: Formula E — same as sagittal_vertical_axis.
        if knee is not None and ankle is not None and foot is not None:
            shank = knee - ankle
            foot_seg = foot - ankle
            denom = np.linalg.norm(shank) * np.linalg.norm(foot_seg) + 1e-12
            cos_val = np.clip(np.dot(shank, foot_seg) / denom, -1, 1)
            unsigned = np.degrees(np.arccos(cos_val))
            cross = float(shank[0] * foot_seg[1] - shank[1] * foot_seg[0])
            result[f"ankle_{side}"] = (90.0 - unsigned) if cross >= 0 else -(90.0 - unsigned)
        else:
            result[f"ankle_{side}"] = np.nan

    return result


# ── Ankle projection-t correction ────────────────────────────────────
#
# The MediaPipe ANKLE landmark is a surface feature (lateral or medial
# malleolus) rather than the true joint centre.  During gait — especially
# during loading response — the landmark "slides" along the foot line
# (HEEL → FOOT_INDEX), which distorts the computed ankle angle.
#
# We quantify this sliding via the *projection parameter t*: the
# normalised position of ANKLE projected onto the HEEL→FOOT_INDEX line.
#   t = dot(ANKLE − HEEL, foot_dir) / |foot_dir|²
# When t drops below its median the landmark has slid toward the heel
# and the ankle angle contains an artefact.
#
# The correction estimates the angle error as a linear function of the
# deviation Δt = t_median − t_frame, calibrated on flat-foot phases
# (|HEEL.y − FOOT.y| < threshold) where the true ankle angle is nearly
# constant.  Any ankle-angle variation during flat-foot is attributed to
# landmark sliding and regressed on Δt.
#
# Reference: internal validation against Vicon (trial 13, R² > 0.85).


def _ankle_projection_t(
    ankle: np.ndarray,
    heel: np.ndarray,
    foot: np.ndarray,
) -> Optional[float]:
    """Normalised projection of ANKLE onto the HEEL→FOOT line.

    Returns the scalar *t* such that the projection of ANKLE onto the
    directed segment HEEL→FOOT is ``HEEL + t * |HF| * unit(HF)``.
    Returns ``None`` when HEEL and FOOT overlap.
    """
    fd = foot - heel
    d_hf = np.linalg.norm(fd)
    if d_hf < 1e-6:
        return None
    return float(np.dot(ankle - heel, fd) / (d_hf * d_hf))


def _correct_ankle_projection(
    frames: list,
    angle_frames: list,
    model: str,
    flat_foot_threshold: float = 0.012,
) -> None:
    """Apply projection-t correction to ankle angles in-place.

    For each side (L, R):
      1. Compute projection t for every frame.
      2. Identify flat-foot phases (|HEEL.y − FOOT.y| < threshold).
      3. During flat-foot, regress ankle-angle variation on Δt = t_med − t.
      4. Subtract the estimated artefact from ALL frames.

    Parameters
    ----------
    frames : list[dict]
        Raw frames from the pivot JSON (with ``landmarks``).
    angle_frames : list[dict]
        Computed angle dicts (modified in-place).
    model : str
        Extraction model name (used to decide foot-landmark estimation).
    flat_foot_threshold : float
        Maximum |HEEL.y − FOOT.y| to consider the foot as flat (in
        normalised image coordinates).  Default 0.012.
    """
    needs_foot = model != "mediapipe"

    for side, prefix in [("L", "LEFT"), ("R", "RIGHT")]:
        key = f"ankle_{side}"

        # --- Pass 1: collect t values and raw ankle angles ---------------
        t_values = []
        raw_angles = []
        flat_mask = []

        for i, frame in enumerate(frames):
            f = _estimate_foot_landmarks(frame) if needs_foot else frame
            knee = _get_xy(f, f"{prefix}_KNEE")
            ankle = _get_xy(f, f"{prefix}_ANKLE")
            heel = _get_xy(f, f"{prefix}_HEEL")
            foot = _get_foot_index_from_toes(f, prefix)

            if any(v is None for v in [knee, ankle, heel, foot]):
                t_values.append(np.nan)
                raw_angles.append(np.nan)
                flat_mask.append(False)
                continue

            t = _ankle_projection_t(ankle, heel, foot)
            t_values.append(t if t is not None else np.nan)

            a = angle_frames[i].get(key)
            raw_angles.append(a if a is not None and not np.isnan(a) else np.nan)

            # Flat-foot detection
            is_flat = abs(heel[1] - foot[1]) < flat_foot_threshold
            flat_mask.append(is_flat)

        t_values = np.array(t_values, dtype=float)
        raw_angles = np.array(raw_angles, dtype=float)
        flat_mask = np.array(flat_mask, dtype=bool)

        # Median projection parameter
        valid_t = t_values[~np.isnan(t_values)]
        if len(valid_t) < 10:
            logger.debug("ankle projection correction %s: too few t values", side)
            continue
        t_med = float(np.median(valid_t))

        # --- Pass 2: calibrate on flat-foot phases -----------------------
        cal_mask = (
            flat_mask
            & ~np.isnan(t_values)
            & ~np.isnan(raw_angles)
        )
        if np.sum(cal_mask) < 10:
            logger.debug(
                "ankle projection correction %s: only %d flat-foot frames, skipping",
                side, int(np.sum(cal_mask)),
            )
            continue

        dt_flat = t_med - t_values[cal_mask]
        angles_flat = raw_angles[cal_mask]
        mean_flat = float(np.mean(angles_flat))
        variation_flat = angles_flat - mean_flat

        # Linear regression: variation = k * Δt  (force intercept ≈ 0)
        if np.std(dt_flat) < 1e-8:
            continue
        k = float(np.polyfit(dt_flat, variation_flat, 1)[0])

        logger.info(
            "ankle projection correction %s: t_med=%.3f, k=%.1f, "
            "n_flat=%d",
            side, t_med, k, int(np.sum(cal_mask)),
        )

        # --- Pass 3: apply correction ------------------------------------
        for i in range(len(angle_frames)):
            a = angle_frames[i].get(key)
            if a is None or np.isnan(a) or np.isnan(t_values[i]):
                continue
            dt = t_med - t_values[i]
            angle_frames[i][key] = float(a - k * dt)


# ── Method registry ──────────────────────────────────────────────────

ANGLE_METHODS: Dict[str, Callable] = {
    "sagittal_vertical_axis": _method_sagittal_vertical_axis,
    "sagittal_classic": _method_sagittal_classic,
}


def register_angle_method(name: str, func: Callable):
    """Register a custom angle computation method.

    The function must accept (frame: dict, model: str) -> dict
    and return a dict with keys: frame_idx, hip_L, hip_R, knee_L, knee_R,
    ankle_L, ankle_R, trunk_angle, pelvis_tilt, landmark_positions.
    """
    ANGLE_METHODS[name] = func


def list_angle_methods() -> list:
    """Return available angle computation method names."""
    return list(ANGLE_METHODS.keys())


# ── Public API ───────────────────────────────────────────────────────


def compute_angles(
    data: dict,
    method: str = "sagittal_vertical_axis",
    correction_factor: float = 0.8,
    calibrate: bool = True,
    calibration_frames: int = 30,
    calibration_joints: Optional[list] = None,
    min_confidence: float = 0.0,
    correct_ankle_sliding: bool = True,
) -> dict:
    """Compute joint angles and add to pivot JSON.

    Parameters
    ----------
    data : dict
        Pivot JSON dict with ``frames`` populated.
    method : str, optional
        Angle computation method (default ``"sagittal_vertical_axis"``).
        See ``list_angle_methods()`` for available methods.
    correction_factor : float, optional
        Scale factor for 2D ROM overestimation (default 0.8 for
        MediaPipe, use 1.0 for 3D models).
    calibrate : bool, optional
        Subtract neutral angle from the first N frames (default True).
    calibration_frames : int, optional
        Number of frames for neutral calibration (default 30).
    calibration_joints : list of str, optional
        Joints to calibrate (default: ``["ankle_L", "ankle_R"]``).
    min_confidence : float, optional
        Skip angle computation on frames with confidence below this
        threshold (default 0.0, i.e. compute on all frames).  Skipped
        frames still appear in the output with all joint values set to
        ``None`` so that frame indices stay aligned.
    correct_ankle_sliding : bool, optional
        Apply projection-t correction to compensate for the ANKLE
        landmark sliding along the foot line during gait (default
        ``True``).  The correction is self-calibrated using flat-foot
        phases and requires no external reference data.

    Returns
    -------
    dict
        Modified *data* dict with ``angles`` populated.

    Raises
    ------
    ValueError
        If *data* has no frames or *method* is unknown.
    """
    if not data.get("frames"):
        raise ValueError("No frames in data. Run extract() first.")

    if method not in ANGLE_METHODS:
        available = ", ".join(ANGLE_METHODS.keys())
        raise ValueError(f"Unknown angle method '{method}'. Available: {available}")

    if calibration_joints is None:
        calibration_joints = ["ankle_L", "ankle_R"]

    method_func = ANGLE_METHODS[method]
    model = data.get("extraction", {}).get("model", "mediapipe")

    # Detect walking direction (only relevant for sagittal_vertical_axis)
    walking_direction = _detect_walking_direction(data)
    logger.info("Detected walking direction: %s", walking_direction)

    # Compute per-frame angles (skip low-confidence frames)
    joint_keys = ["hip_L", "hip_R", "knee_L", "knee_R", "ankle_L", "ankle_R"]
    angle_frames = []
    n_skipped = 0
    for frame in data["frames"]:
        conf = frame.get("confidence", 0.0)
        if conf < min_confidence:
            # Placeholder with None values so frame indices stay aligned
            af = {"frame_idx": frame["frame_idx"]}
            for k in joint_keys:
                af[k] = np.nan
            af["trunk_angle"] = np.nan
            af["pelvis_tilt"] = np.nan
            af["landmark_positions"] = {}
            angle_frames.append(af)
            n_skipped += 1
        else:
            angle_frames.append(method_func(frame, model))
    if n_skipped:
        logger.info("Skipped %d/%d low-confidence frames (< %.2f)",
                     n_skipped, len(data["frames"]), min_confidence)

    # If the subject walks right-to-left the sagittal_vertical_axis
    # hip sign is mirrored.  Invert so that flexion stays positive.
    if walking_direction == "right_to_left" and method == "sagittal_vertical_axis":
        for af in angle_frames:
            for side in ("L", "R"):
                key = f"hip_{side}"
                v = af.get(key)
                if v is not None and not np.isnan(v):
                    af[key] = -v

    # Trunk angle also needs direction correction for right-to-left walking.
    if walking_direction == "right_to_left":
        for af in angle_frames:
            v = af.get("trunk_angle")
            if v is not None and not np.isnan(v):
                af["trunk_angle"] = -v

    # Ankle landmark sliding correction (projection-t method)
    if correct_ankle_sliding:
        _correct_ankle_projection(data["frames"], angle_frames, model)

    # Apply correction factor
    if correction_factor != 1.0:
        joint_keys = ["hip_L", "hip_R", "knee_L", "knee_R", "ankle_L", "ankle_R"]
        for af in angle_frames:
            for key in joint_keys:
                v = af.get(key)
                if v is not None and not np.isnan(v):
                    af[key] = v * correction_factor

    # Normalize hip angles to [-180, 180] (no unwrap — gait is cyclic)
    for side in ("L", "R"):
        key = f"hip_{side}"
        for af in angle_frames:
            v = af.get(key)
            if v is not None and not np.isnan(v):
                af[key] = float(((v + 180) % 360) - 180)

    # Unwrap pelvis tilt to remove ±180° discontinuities
    pelvis_vals = [af.get("pelvis_tilt") for af in angle_frames]
    pelvis_vals = [v if v is not None else np.nan for v in pelvis_vals]
    pelvis_unwrapped = _unwrap_angles(pelvis_vals)
    for af, v in zip(angle_frames, pelvis_unwrapped):
        af["pelvis_tilt"] = v

    # Neutral calibration
    if calibrate and len(angle_frames) >= calibration_frames:
        for key in calibration_joints:
            vals = [af[key] for af in angle_frames[:calibration_frames]
                    if not np.isnan(af.get(key, np.nan))]
            if vals:
                offset = float(np.median(vals))
                for af in angle_frames:
                    v = af.get(key)
                    if v is not None and not np.isnan(v):
                        af[key] = v - offset

    # Convert NaN to None for JSON serialization
    for af in angle_frames:
        for key in list(af.keys()):
            if key == "landmark_positions":
                continue
            if isinstance(af[key], float) and np.isnan(af[key]):
                af[key] = None

    data["angles"] = {
        "method": method,
        "correction_factor": correction_factor,
        "calibrated": calibrate,
        "calibration_frames": calibration_frames if calibrate else None,
        "calibration_joints": calibration_joints if calibrate else [],
        "ankle_sliding_correction": correct_ankle_sliding,
        "walking_direction": walking_direction,
        "joints": ["hip_L", "hip_R", "knee_L", "knee_R", "ankle_L", "ankle_R"],
        "extra": ["trunk_angle", "pelvis_tilt"],
        "frames": angle_frames,
    }

    return data


# ── Ankle swap detection (dual-method) ───────────────────────────────
#
# Sapiens can swap LEFT_ANKLE / RIGHT_ANKLE labels between frames while
# HEEL and FOOT_INDEX remain correctly labelled.  We exploit this by
# computing the ankle dorsiflexion angle with two independent methods:
#
#   Method A — uses ANKLE as the angle vertex (affected by swap)
#   Method B — uses HEEL  as the angle vertex (immune to swap)
#
# When the two methods diverge beyond a threshold the ANKLE label is
# deemed swapped and the angle from Method B is used instead.
#
# Reference for the clinical goniometric model:
#   Norkin CC, White DJ. Measurement of Joint Motion: A Guide to
#   Goniometry. 5th ed. F.A. Davis; 2016. Chapter 11 (Ankle).


def ankle_angle_method_A(
    knee: np.ndarray,
    ankle: np.ndarray,
    heel: np.ndarray,
    foot_index: np.ndarray,
) -> float:
    """Ankle dorsiflexion via the shank-foot angle (shank origin = ANKLE).

    Replicates the same geometry used by ``_method_sagittal_vertical_axis``
    in this module:

      - shank vector = KNEE - ANKLE  (points proximally up the leg)
      - foot vector  = FOOT_INDEX - HEEL  (points distally along the foot)

    The foot vector uses HEEL -> FOOT_INDEX (not ANKLE -> FOOT_INDEX)
    because both points lie on the same rigid segment and give a more
    stable orientation.  The ANKLE dependency comes only from the shank
    direction vector.

    At anatomical neutral the angle between shank and foot is ~90 deg.
    Dorsiflexion (foot up) decreases this angle; plantarflexion (foot
    down) increases it.

    Formula
    -------
    unsigned = arccos(dot(shank, foot) / (|shank| * |foot|))
    cross    = shank.x * foot.y - shank.y * foot.x
    angle    = (90 - unsigned) if cross >= 0 else -(90 - unsigned)

    Convention: dorsiflexion > 0, plantarflexion < 0, neutral = 0 deg.

    Why it is affected by a swap
    ----------------------------
    If the ANKLE landmark belongs to the contralateral limb while HEEL
    and FOOT_INDEX are correct, the shank direction vector is wrong and
    the computed angle becomes meaningless.

    Parameters
    ----------
    knee, ankle, heel, foot_index : np.ndarray
        (x, y) coordinates in image space.

    Returns
    -------
    float
        Ankle angle in degrees (dorsiflexion positive).
    """
    shank = knee - ankle        # points up along the leg
    foot_seg = foot_index - heel  # points forward along the foot

    n1 = np.linalg.norm(shank)
    n2 = np.linalg.norm(foot_seg)
    if n1 < 1e-10 or n2 < 1e-10:
        return np.nan

    cos_a = np.clip(np.dot(shank, foot_seg) / (n1 * n2), -1.0, 1.0)
    unsigned = np.degrees(np.arccos(cos_a))
    cross = float(shank[0] * foot_seg[1] - shank[1] * foot_seg[0])

    return (90.0 - unsigned) if cross >= 0 else -(90.0 - unsigned)


def ankle_angle_method_B(
    knee: np.ndarray,
    heel: np.ndarray,
    foot_index: np.ndarray,
) -> float:
    """Ankle dorsiflexion via the heel-pivot method (does NOT use ANKLE).

    The HEEL landmark sits directly below the malleolus in the sagittal
    plane, so KNEE -> HEEL is a valid proxy for the tibial axis.  The
    foot axis is HEEL -> FOOT_INDEX (identical to Method A's foot vector).

    This is biomechanically equivalent to measuring the angle between
    the leg and the foot with the goniometer pivot placed on the
    calcaneus instead of the malleolus.  In 2D sagittal projection the
    angular difference is very small because HEEL is only a few cm below
    ANKLE.

    Vectors
    -------
    shank_proxy = KNEE - HEEL      (tibial axis proxy, points up)
    foot_seg    = FOOT_INDEX - HEEL (foot axis, points forward)

    Formula (identical to Method A, just different shank origin)
    -------
    unsigned = arccos(dot(shank_proxy, foot_seg) / (|shank| * |foot|))
    cross    = shank_proxy.x * foot_seg.y - shank_proxy.y * foot_seg.x
    angle    = (90 - unsigned) if cross >= 0 else -(90 - unsigned)

    Convention: same as Method A.

    Why it is immune to the swap
    ----------------------------
    This method uses only KNEE, HEEL, and FOOT_INDEX.  Since the swap
    problem only affects ANKLE, Method B always returns the correct
    angle for the ipsilateral limb regardless of ANKLE labelling.

    Parameters
    ----------
    knee, heel, foot_index : np.ndarray
        (x, y) coordinates in image space.

    Returns
    -------
    float
        Ankle angle in degrees (dorsiflexion positive).
    """
    shank = knee - heel           # tibial axis proxy, points up
    foot_seg = foot_index - heel  # foot axis, points forward

    n1 = np.linalg.norm(shank)
    n2 = np.linalg.norm(foot_seg)
    if n1 < 1e-10 or n2 < 1e-10:
        return np.nan

    cos_a = np.clip(np.dot(shank, foot_seg) / (n1 * n2), -1.0, 1.0)
    unsigned = np.degrees(np.arccos(cos_a))
    cross = float(shank[0] * foot_seg[1] - shank[1] * foot_seg[0])

    return (90.0 - unsigned) if cross >= 0 else -(90.0 - unsigned)


def detect_ankle_swap(
    frame: dict,
    side: str,
    threshold_deg: float = 8.0,
) -> dict:
    """Detect whether the ANKLE label is swapped for *side* in *frame*.

    Computes Method A (uses ANKLE) and Method B (uses HEEL as pivot).
    If ``|A - B| > threshold_deg`` the ANKLE label is flagged as swapped,
    and the Method B angle is returned as the corrected value.

    The default threshold of 8 degrees was chosen as follows:

    * In normal conditions HEEL is ~2-4 cm below ANKLE in sagittal view.
      For a typical shank length of 40 cm this creates an angular offset
      of arctan(0.04/0.40) = 5.7 deg at most between methods A and B.
    * Markerless pose noise adds ~2-3 deg RMS per landmark.
    * A swapped ANKLE typically produces a discrepancy of 15-40+ deg
      because the contralateral malleolus is at a completely different
      position.
    * Therefore 8 deg provides comfortable separation: normal noise
      stays below 6 deg, while a true swap exceeds 15 deg.

    Parameters
    ----------
    frame : dict
        Single frame dict with ``landmarks``.
    side : str
        ``"LEFT"`` or ``"RIGHT"``.
    threshold_deg : float
        Absolute difference in degrees above which ANKLE is deemed
        swapped (default 8.0).

    Returns
    -------
    dict
        Keys:

        - ``swapped`` (bool): True if swap detected.
        - ``angle_method_A`` (float): Angle from Method A (uses ANKLE).
        - ``angle_method_B`` (float): Angle from Method B (HEEL pivot).
        - ``delta_deg`` (float): ``|A - B|``.
        - ``corrected_angle`` (float): Best estimate of the true angle
          (Method B if swapped, Method A otherwise).
    """
    knee = _get_xy(frame, f"{side}_KNEE")
    ankle = _get_xy(frame, f"{side}_ANKLE")
    heel = _get_xy(frame, f"{side}_HEEL")
    foot = _get_foot_index_from_toes(frame, side)

    result = {
        "swapped": False,
        "angle_method_A": np.nan,
        "angle_method_B": np.nan,
        "delta_deg": np.nan,
        "corrected_angle": np.nan,
    }

    if knee is None or heel is None or foot is None:
        return result

    # Method B is always computable (does not need ANKLE)
    angle_B = ankle_angle_method_B(knee, heel, foot)
    result["angle_method_B"] = float(angle_B)

    if ankle is None:
        # ANKLE missing entirely — use Method B
        result["corrected_angle"] = float(angle_B)
        return result

    angle_A = ankle_angle_method_A(knee, ankle, heel, foot)
    result["angle_method_A"] = float(angle_A)

    if np.isnan(angle_A) or np.isnan(angle_B):
        result["corrected_angle"] = float(angle_B) if not np.isnan(angle_B) else float(angle_A)
        return result

    delta = abs(angle_A - angle_B)
    result["delta_deg"] = float(delta)

    if delta > threshold_deg:
        result["swapped"] = True
        result["corrected_angle"] = float(angle_B)
        logger.warning(
            "Ankle swap detected on %s (frame %s): "
            "Method A=%.1f deg, Method B=%.1f deg, delta=%.1f deg > %.1f deg threshold. "
            "Using Method B (HEEL pivot) as corrected value.",
            side,
            frame.get("frame_idx", "?"),
            angle_A,
            angle_B,
            delta,
            threshold_deg,
        )
    else:
        result["corrected_angle"] = float(angle_A)

    return result


def correct_ankle_swaps(
    data: dict,
    threshold_deg: float = 8.0,
) -> dict:
    """Scan all frames and correct ankle angles where swaps are detected.

    This function should be called AFTER ``compute_angles()`` (which
    populates ``data["angles"]["frames"]``).  It re-examines each frame
    using the dual-method approach and patches ``ankle_L`` / ``ankle_R``
    when a swap is detected.

    Also stores diagnostic metadata in ``data["angles"]["ankle_swap"]``.

    Parameters
    ----------
    data : dict
        Pivot JSON dict with ``frames`` and ``angles`` populated.
    threshold_deg : float
        Threshold for swap detection (default 8.0 degrees).

    Returns
    -------
    dict
        Modified *data* with corrected ankle angles.
    """
    pose_frames = data.get("frames", [])
    angle_frames = data.get("angles", {}).get("frames", [])

    if not pose_frames or not angle_frames:
        return data

    n_swaps_L = 0
    n_swaps_R = 0
    swap_frames = []

    for i, (pf, af) in enumerate(zip(pose_frames, angle_frames)):
        for side, key in [("LEFT", "ankle_L"), ("RIGHT", "ankle_R")]:
            result = detect_ankle_swap(pf, side, threshold_deg)

            if result["swapped"]:
                # Replace the angle with the corrected value
                af[key] = result["corrected_angle"]
                swap_frames.append({
                    "frame_idx": pf.get("frame_idx", i),
                    "side": side,
                    "angle_A": result["angle_method_A"],
                    "angle_B": result["angle_method_B"],
                    "delta_deg": result["delta_deg"],
                })
                if side == "LEFT":
                    n_swaps_L += 1
                else:
                    n_swaps_R += 1

    total = len(pose_frames)
    data["angles"]["ankle_swap"] = {
        "threshold_deg": threshold_deg,
        "n_swaps_left": n_swaps_L,
        "n_swaps_right": n_swaps_R,
        "pct_swapped_left": round(100.0 * n_swaps_L / total, 1) if total else 0.0,
        "pct_swapped_right": round(100.0 * n_swaps_R / total, 1) if total else 0.0,
        "swap_frames": swap_frames,
    }

    if n_swaps_L + n_swaps_R > 0:
        logger.info(
            "Ankle swap correction: %d left, %d right swaps corrected "
            "out of %d frames (threshold=%.1f deg).",
            n_swaps_L, n_swaps_R, total, threshold_deg,
        )

    return data


# ── Extended angle functions ─────────────────────────────────────────


def _head_angle(frame: dict) -> float:
    """Compute head posture angle (head-trunk angle in the sagittal plane).

    Measures the angle between the head vector (ear_midpoint to nose) and
    the trunk vector (hip_center to shoulder_center). A positive value
    indicates forward head posture.

    Uses landmarks: NOSE, LEFT_EAR, RIGHT_EAR, LEFT_SHOULDER,
    RIGHT_SHOULDER, LEFT_HIP, RIGHT_HIP.

    Args:
        frame: Single frame dict with ``landmarks``.

    Returns:
        Head angle in degrees. Positive = forward head posture.
        Returns NaN if required landmarks are missing.
    """
    nose = _get_xy(frame, "NOSE")
    l_ear = _get_xy(frame, "LEFT_EAR")
    r_ear = _get_xy(frame, "RIGHT_EAR")
    l_shoulder = _get_xy(frame, "LEFT_SHOULDER")
    r_shoulder = _get_xy(frame, "RIGHT_SHOULDER")
    l_hip = _get_xy(frame, "LEFT_HIP")
    r_hip = _get_xy(frame, "RIGHT_HIP")

    if any(v is None for v in [nose, l_ear, r_ear, l_shoulder, r_shoulder, l_hip, r_hip]):
        return np.nan

    ear_mid = (l_ear + r_ear) / 2
    shoulder_center = (l_shoulder + r_shoulder) / 2
    hip_center = (l_hip + r_hip) / 2

    # Head vector: ear midpoint to nose
    head_vec = nose - ear_mid
    # Trunk vector: hip center to shoulder center (points upward)
    trunk_vec = shoulder_center - hip_center

    angle = _angle_between(head_vec, trunk_vec)

    # Sign convention: positive = nose forward of trunk line
    # Use cross product sign to determine direction
    cross = head_vec[0] * trunk_vec[1] - head_vec[1] * trunk_vec[0]
    return angle if cross >= 0 else -angle


def _arm_angles(frame: dict) -> dict:
    """Compute shoulder flexion and elbow flexion angles for both sides.

    Shoulder flexion: angle of the upper arm vector relative to trunk
    vertical (hip_center to shoulder_center direction). 0 = arm hanging
    straight down.

    Elbow flexion: 180 minus the interior angle at the elbow
    (shoulder-elbow-wrist). 0 = full extension, positive = flexion.

    Uses landmarks: LEFT/RIGHT SHOULDER, ELBOW, WRIST, HIP.

    Args:
        frame: Single frame dict with ``landmarks``.

    Returns:
        Dict with keys: shoulder_flex_L, shoulder_flex_R,
        elbow_flex_L, elbow_flex_R (all in degrees, NaN if missing).
    """
    l_shoulder = _get_xy(frame, "LEFT_SHOULDER")
    r_shoulder = _get_xy(frame, "RIGHT_SHOULDER")
    l_hip = _get_xy(frame, "LEFT_HIP")
    r_hip = _get_xy(frame, "RIGHT_HIP")

    result = {
        "shoulder_flex_L": np.nan,
        "shoulder_flex_R": np.nan,
        "elbow_flex_L": np.nan,
        "elbow_flex_R": np.nan,
    }

    # Trunk vertical direction (used for shoulder flexion reference)
    trunk_vertical = None
    if (l_shoulder is not None and r_shoulder is not None
            and l_hip is not None and r_hip is not None):
        shoulder_center = (l_shoulder + r_shoulder) / 2
        hip_center = (l_hip + r_hip) / 2
        # Points downward in image coords (hip below shoulder)
        trunk_vertical = hip_center - shoulder_center

    for side, prefix in [("L", "LEFT"), ("R", "RIGHT")]:
        shoulder = _get_xy(frame, f"{prefix}_SHOULDER")
        elbow = _get_xy(frame, f"{prefix}_ELBOW")
        wrist = _get_xy(frame, f"{prefix}_WRIST")

        # Shoulder flexion
        if shoulder is not None and elbow is not None and trunk_vertical is not None:
            upper_arm = elbow - shoulder
            result[f"shoulder_flex_{side}"] = _angle_between(upper_arm, trunk_vertical)

        # Elbow flexion
        if shoulder is not None and elbow is not None and wrist is not None:
            v1 = shoulder - elbow
            v2 = wrist - elbow
            result[f"elbow_flex_{side}"] = 180.0 - _angle_between(v1, v2)

    return result


def _pelvis_sagittal_tilt(frame: dict) -> float:
    """Compute pelvis anterior/posterior tilt in the sagittal plane.

    Measures the angle between the hip_center-to-shoulder_center vector
    and the vertical axis. This differs from the existing ``_pelvis_tilt``
    which measures frontal-plane (lateral) tilt.

    Positive = anterior tilt (trunk leaning forward relative to hips).

    Uses landmarks: LEFT_HIP, RIGHT_HIP, LEFT_SHOULDER, RIGHT_SHOULDER.

    Args:
        frame: Single frame dict with ``landmarks``.

    Returns:
        Sagittal pelvis tilt in degrees. Positive = anterior tilt.
        Returns NaN if required landmarks are missing.
    """
    l_hip = _get_xy(frame, "LEFT_HIP")
    r_hip = _get_xy(frame, "RIGHT_HIP")
    l_shoulder = _get_xy(frame, "LEFT_SHOULDER")
    r_shoulder = _get_xy(frame, "RIGHT_SHOULDER")

    if any(v is None for v in [l_hip, r_hip, l_shoulder, r_shoulder]):
        return np.nan

    hip_center = (l_hip + r_hip) / 2
    shoulder_center = (l_shoulder + r_shoulder) / 2

    # Vector from hip center to shoulder center
    trunk_vec = shoulder_center - hip_center
    # Vertical (pointing up in image = negative y)
    vertical = np.array([0.0, -1.0])

    angle = _angle_between(trunk_vec, vertical)

    # Sign: positive when trunk tilts forward (shoulder x > hip x)
    return angle if trunk_vec[0] >= 0 else -angle


def _depth_enhanced_angles(frame: dict) -> Optional[dict]:
    """Correct 2D joint angles using depth data to compensate for foreshortening.

    When a limb segment is angled toward or away from the camera, its
    projected 2D length is shorter than the true 3D length. If depth data
    is available (from Sapiens depth estimation), we can estimate the
    foreshortening factor and correct the 2D angles.

    The correction uses the depth difference between proximal and distal
    joints to estimate the out-of-plane angle, then scales the 2D angle
    by the inverse cosine of that angle.

    Only works when the frame has a ``landmark_depths`` key.

    Args:
        frame: Single frame dict with ``landmarks`` and optionally
            ``landmark_depths`` mapping landmark names to depth values.

    Returns:
        Dict with corrected angles (hip_L, hip_R, knee_L, knee_R,
        ankle_L, ankle_R) or None if no depth data is available.
    """
    depths = frame.get("landmark_depths")
    if depths is None:
        return None

    result = {}

    for side, prefix in [("L", "LEFT"), ("R", "RIGHT")]:
        # Hip angle correction: use hip-knee depth difference
        hip = _get_xy(frame, f"{prefix}_HIP")
        knee = _get_xy(frame, f"{prefix}_KNEE")
        ankle = _get_xy(frame, f"{prefix}_ANKLE")

        d_hip = depths.get(f"{prefix}_HIP")
        d_knee = depths.get(f"{prefix}_KNEE")
        d_ankle = depths.get(f"{prefix}_ANKLE")

        # Hip correction
        if (hip is not None and knee is not None
                and d_hip is not None and d_knee is not None):
            seg_2d = np.linalg.norm(knee - hip)
            depth_diff = abs(d_knee - d_hip)
            if seg_2d > 1e-6:
                # Estimate true 3D length
                seg_3d = np.sqrt(seg_2d ** 2 + depth_diff ** 2)
                correction = seg_3d / seg_2d if seg_2d > 0 else 1.0
                # Simple scaling: multiply 2D angle by correction factor
                # (capped to avoid extreme corrections)
                correction = min(correction, 2.0)
            else:
                correction = 1.0
            result[f"hip_{side}_correction"] = float(correction)
        else:
            result[f"hip_{side}_correction"] = 1.0

        # Knee correction
        if (knee is not None and ankle is not None
                and d_knee is not None and d_ankle is not None):
            seg_2d = np.linalg.norm(ankle - knee)
            depth_diff = abs(d_ankle - d_knee)
            if seg_2d > 1e-6:
                seg_3d = np.sqrt(seg_2d ** 2 + depth_diff ** 2)
                correction = min(seg_3d / seg_2d, 2.0)
            else:
                correction = 1.0
            result[f"knee_{side}_correction"] = float(correction)
        else:
            result[f"knee_{side}_correction"] = 1.0

        # Ankle correction
        if (ankle is not None and d_ankle is not None):
            foot = _get_xy(frame, f"{prefix}_FOOT_INDEX")
            d_foot = depths.get(f"{prefix}_FOOT_INDEX")
            if foot is not None and d_foot is not None:
                seg_2d = np.linalg.norm(foot - ankle)
                depth_diff = abs(d_foot - d_ankle)
                if seg_2d > 1e-6:
                    seg_3d = np.sqrt(seg_2d ** 2 + depth_diff ** 2)
                    correction = min(seg_3d / seg_2d, 2.0)
                else:
                    correction = 1.0
                result[f"ankle_{side}_correction"] = float(correction)
            else:
                result[f"ankle_{side}_correction"] = 1.0
        else:
            result[f"ankle_{side}_correction"] = 1.0

    return result


def compute_frontal_angles(data: dict) -> dict:
    """Compute frontal-plane angles from depth-enhanced pose data.

    Requires ``landmark_depths`` in frame data (from Sapiens depth
    estimation). Computes hip abduction and knee valgus/varus for
    each side.

    Hip abduction: angle of the thigh in the frontal plane relative
    to vertical, using the mediolateral (x) and depth (z) components.

    Knee valgus: frontal-plane angle at the knee using hip, knee, ankle
    positions in the x-z plane.

    Results are stored in ``data["angles_frontal"]``.

    Args:
        data: Pivot JSON dict with ``frames`` populated, ideally with
            ``landmark_depths`` in each frame.

    Returns:
        Modified *data* dict with ``angles_frontal`` populated.
        If no depth data is available, ``data["angles_frontal"]`` is set
        to None.
    """
    frames = data.get("frames", [])
    if not frames:
        data["angles_frontal"] = None
        return data

    # Check if any frame has depth data
    has_depth = any(f.get("landmark_depths") is not None for f in frames)
    if not has_depth:
        data["angles_frontal"] = None
        return data

    frontal_frames = []
    for frame in frames:
        depths = frame.get("landmark_depths", {})
        if not depths:
            frontal_frames.append({
                "frame_idx": frame["frame_idx"],
                "hip_abduction_L": None,
                "hip_abduction_R": None,
                "knee_valgus_L": None,
                "knee_valgus_R": None,
            })
            continue

        result = {"frame_idx": frame["frame_idx"]}

        for side, prefix in [("L", "LEFT"), ("R", "RIGHT")]:
            hip = _get_xy(frame, f"{prefix}_HIP")
            knee = _get_xy(frame, f"{prefix}_KNEE")
            ankle = _get_xy(frame, f"{prefix}_ANKLE")

            d_hip = depths.get(f"{prefix}_HIP")
            d_knee = depths.get(f"{prefix}_KNEE")
            d_ankle = depths.get(f"{prefix}_ANKLE")

            # Hip abduction: frontal plane angle of thigh
            if (hip is not None and knee is not None
                    and d_hip is not None and d_knee is not None):
                # Frontal plane: x (mediolateral) and z (depth)
                dx = knee[0] - hip[0]
                dz = d_knee - d_hip
                # Vertical in frontal plane is pure z-depth direction
                frontal_vec = np.array([dx, dz])
                vertical = np.array([0.0, 1.0])
                angle = _angle_between(frontal_vec, vertical)
                result[f"hip_abduction_{side}"] = float(angle)
            else:
                result[f"hip_abduction_{side}"] = None

            # Knee valgus: frontal plane angle at knee
            if (hip is not None and knee is not None and ankle is not None
                    and d_hip is not None and d_knee is not None and d_ankle is not None):
                # Thigh in frontal plane
                v1_x = hip[0] - knee[0]
                v1_z = d_hip - d_knee
                # Shank in frontal plane
                v2_x = ankle[0] - knee[0]
                v2_z = d_ankle - d_knee
                v1 = np.array([v1_x, v1_z])
                v2 = np.array([v2_x, v2_z])
                angle = 180.0 - _angle_between(v1, v2)
                result[f"knee_valgus_{side}"] = float(angle)
            else:
                result[f"knee_valgus_{side}"] = None

        frontal_frames.append(result)

    data["angles_frontal"] = {
        "frames": frontal_frames,
        "joints": ["hip_abduction_L", "hip_abduction_R", "knee_valgus_L", "knee_valgus_R"],
    }
    return data


def compute_extended_angles(data: dict) -> dict:
    """Compute extended angles (head, arms, pelvis sagittal) and add to angle frames.

    Call AFTER ``compute_angles()``. Adds the following keys to each
    angle frame in ``data["angles"]["frames"]``:

    - ``head_angle``: head posture angle (forward head = positive)
    - ``shoulder_flex_L``, ``shoulder_flex_R``: shoulder flexion
    - ``elbow_flex_L``, ``elbow_flex_R``: elbow flexion
    - ``pelvis_sagittal_tilt``: anterior/posterior pelvic tilt

    Also attempts depth-enhanced angle correction if depth data is
    available in the pose frames, storing correction factors in each
    angle frame.

    Args:
        data: Pivot JSON dict with ``angles`` populated (via
            ``compute_angles()``).

    Returns:
        Modified *data* dict with extended angles added.

    Raises:
        ValueError: If ``data["angles"]`` is not populated.
    """
    if not data.get("angles") or not data["angles"].get("frames"):
        raise ValueError("No angles computed. Run compute_angles() first.")

    angle_frames = data["angles"]["frames"]
    pose_frames = data.get("frames", [])

    for i, af in enumerate(angle_frames):
        if i >= len(pose_frames):
            break
        frame = pose_frames[i]

        # Head angle
        af["head_angle"] = _head_angle(frame)

        # Arm angles
        arm = _arm_angles(frame)
        af["shoulder_flex_L"] = arm["shoulder_flex_L"]
        af["shoulder_flex_R"] = arm["shoulder_flex_R"]
        af["elbow_flex_L"] = arm["elbow_flex_L"]
        af["elbow_flex_R"] = arm["elbow_flex_R"]

        # Pelvis sagittal tilt
        af["pelvis_sagittal_tilt"] = _pelvis_sagittal_tilt(frame)

        # Depth-enhanced correction
        depth_corr = _depth_enhanced_angles(frame)
        if depth_corr is not None:
            af["depth_corrections"] = depth_corr

    # Convert NaN to None for JSON serialization in new keys
    extended_keys = [
        "head_angle", "shoulder_flex_L", "shoulder_flex_R",
        "elbow_flex_L", "elbow_flex_R", "pelvis_sagittal_tilt",
    ]
    for af in angle_frames:
        for key in extended_keys:
            val = af.get(key)
            if isinstance(val, float) and np.isnan(val):
                af[key] = None

    return data


def foot_progression_angle(data: dict) -> dict:
    """Compute foot progression angle (heel-to-toe vs horizontal).

    The foot progression angle (FPA) is the angle between the
    heel-to-toe vector and the horizontal axis.  Positive values
    indicate external rotation (out-toeing), negative values indicate
    internal rotation (in-toeing).

    Uses LEFT_HEEL -> LEFT_FOOT_INDEX and RIGHT_HEEL -> RIGHT_FOOT_INDEX.

    Parameters
    ----------
    data : dict
        Pivot JSON dict with ``frames`` populated.

    Returns
    -------
    dict
        ``{"foot_angle_L": [...], "foot_angle_R": [...]}``, where each
        list has one entry per frame (degrees, or ``None`` if landmarks
        are missing).
    """
    frames = data.get("frames", [])
    foot_angle_L = []
    foot_angle_R = []

    for frame in frames:
        for side, result_list in [("LEFT", foot_angle_L), ("RIGHT", foot_angle_R)]:
            heel = _get_xy(frame, f"{side}_HEEL")
            # Prefer real toe midpoint over estimated FOOT_INDEX
            toe = _get_foot_index_from_toes(frame, side)
            if heel is not None and toe is not None:
                vec = toe - heel
                # atan2(y, x) -- but y is inverted in image coords so
                # we use -y to get the real-world angle.  Positive angle
                # means the toe points outward (external rotation).
                angle = np.degrees(np.arctan2(-vec[1], vec[0]))
                result_list.append(float(angle))
            else:
                result_list.append(None)

    return {"foot_angle_L": foot_angle_L, "foot_angle_R": foot_angle_R}
