14_07_expt/preprocessing.py

"""Skeleton cleaning, biomechanical normalization, feature extraction, and resampling."""

from __future__ import annotations

import numpy as np


NUM_JOINTS = 17
NOSE = 0
LEFT_SHOULDER = 5
RIGHT_SHOULDER = 6
LEFT_ELBOW = 7
RIGHT_ELBOW = 8
LEFT_WRIST = 9
RIGHT_WRIST = 10
LEFT_HIP = 11
RIGHT_HIP = 12
LEFT_KNEE = 13
RIGHT_KNEE = 14
LEFT_ANKLE = 15
RIGHT_ANKLE = 16

JOINT_ANGLE_TRIPLETS = (
    (LEFT_SHOULDER, LEFT_ELBOW, LEFT_WRIST),
    (RIGHT_SHOULDER, RIGHT_ELBOW, RIGHT_WRIST),
    (LEFT_ELBOW, LEFT_SHOULDER, LEFT_HIP),
    (RIGHT_ELBOW, RIGHT_SHOULDER, RIGHT_HIP),
    (LEFT_SHOULDER, LEFT_HIP, LEFT_KNEE),
    (RIGHT_SHOULDER, RIGHT_HIP, RIGHT_KNEE),
    (LEFT_HIP, LEFT_KNEE, LEFT_ANKLE),
    (RIGHT_HIP, RIGHT_KNEE, RIGHT_ANKLE),
    (RIGHT_HIP, LEFT_HIP, LEFT_KNEE),
    (LEFT_HIP, RIGHT_HIP, RIGHT_KNEE),
    (LEFT_SHOULDER, NOSE, RIGHT_SHOULDER),
    (LEFT_HIP, NOSE, RIGHT_HIP),
)

BONE_PAIRS = (
    (LEFT_SHOULDER, RIGHT_SHOULDER),
    (LEFT_HIP, RIGHT_HIP),
    (LEFT_SHOULDER, LEFT_ELBOW),
    (LEFT_ELBOW, LEFT_WRIST),
    (RIGHT_SHOULDER, RIGHT_ELBOW),
    (RIGHT_ELBOW, RIGHT_WRIST),
    (LEFT_HIP, LEFT_KNEE),
    (LEFT_KNEE, LEFT_ANKLE),
    (RIGHT_HIP, RIGHT_KNEE),
    (RIGHT_KNEE, RIGHT_ANKLE),
)

KEY_JOINTS = (
    LEFT_SHOULDER,
    RIGHT_SHOULDER,
    LEFT_WRIST,
    RIGHT_WRIST,
    LEFT_HIP,
    RIGHT_HIP,
    LEFT_ANKLE,
    RIGHT_ANKLE,
)


def as_skeleton_array(skeleton: np.ndarray) -> np.ndarray:
    """Convert skeleton data to shape (F, 17, C), where C is 2 or 3."""
    arr = np.asarray(skeleton, dtype=np.float64)
    if arr.ndim == 2 and arr.shape[1] == NUM_JOINTS * 3:
        return arr.reshape(arr.shape[0], NUM_JOINTS, 3)
    if arr.ndim == 2 and arr.shape[1] == NUM_JOINTS * 2:
        return arr.reshape(arr.shape[0], NUM_JOINTS, 2)
    if arr.ndim == 3 and arr.shape[1] == NUM_JOINTS and arr.shape[2] in (2, 3):
        return arr.copy()
    raise ValueError(
        "Expected skeleton shape (F, 51), (F, 34), (F, 17, 3), or (F, 17, 2); "
        f"got {arr.shape}"
    )


def xy(skeleton: np.ndarray) -> np.ndarray:
    """Return only coordinates with shape (F, 17, 2)."""
    return as_skeleton_array(skeleton)[..., :2]


def filter_low_confidence(skeleton: np.ndarray, threshold: float) -> np.ndarray:
    """Mark unreliable joints as NaN.

    Input: (F, 17, 3) or flattened (F, 51). Output: (F, 17, 3).
    Formula: if confidence[f, j] < threshold, set x[f, j] and y[f, j] to NaN.
    Assumption: confidence is in column 2 for each joint.
    """
    arr = as_skeleton_array(skeleton)
    if arr.shape[2] < 3:
        return arr
    out = arr.copy()
    low_confidence = out[..., 2] < threshold
    out[..., :2][low_confidence] = np.nan
    return out


def interpolate_missing(skeleton: np.ndarray) -> np.ndarray:
    """Linearly interpolate NaN coordinate gaps over time.

    Input/output: (F, 17, C). For each joint coordinate, NaNs are replaced with
    np.interp over valid frames. All-NaN tracks are filled with 0.0.
    """
    arr = as_skeleton_array(skeleton)
    out = arr.copy()
    frames = np.arange(out.shape[0])
    for joint in range(out.shape[1]):
        for coord in range(2):
            values = out[:, joint, coord]
            valid = np.isfinite(values)
            if valid.all():
                continue
            if valid.sum() == 0:
                values[:] = 0.0
            elif valid.sum() == 1:
                values[:] = values[valid][0]
            else:
                values[:] = np.interp(frames, frames[valid], values[valid])
            out[:, joint, coord] = values
    return out


def smooth_skeleton(skeleton: np.ndarray, window: int, polyorder: int) -> np.ndarray:
    """Apply Savitzky-Golay smoothing to x/y coordinates only.

    Input/output: (F, 17, C). The filter fits a local polynomial of degree polyorder over
    an odd window and evaluates the smoothed coordinate at each frame. Confidence is unchanged.
    For very short clips, the largest valid odd window is used; if unavailable, input is copied.
    """
    arr = as_skeleton_array(skeleton)
    out = arr.copy()
    n_frames = out.shape[0]
    if n_frames < 3:
        return out

    effective_window = min(int(window), n_frames if n_frames % 2 == 1 else n_frames - 1)
    if effective_window <= polyorder:
        effective_window = polyorder + 2 if (polyorder + 2) % 2 == 1 else polyorder + 3
    if effective_window > n_frames:
        return out
    if effective_window % 2 == 0:
        effective_window -= 1
    if effective_window < 3 or polyorder >= effective_window:
        return out

    try:
        from scipy.signal import savgol_filter

        out[..., :2] = savgol_filter(
            out[..., :2],
            window_length=effective_window,
            polyorder=polyorder,
            axis=0,
            mode="interp",
        )
    except ImportError:
        kernel = np.ones(effective_window, dtype=np.float64) / effective_window
        pad = effective_window // 2
        padded = np.pad(out[..., :2], ((pad, pad), (0, 0), (0, 0)), mode="edge")
        for frame in range(n_frames):
            out[frame, ..., :2] = np.sum(
                padded[frame : frame + effective_window] * kernel[:, None, None], axis=0
            )
    return out


def correct_aspect_ratio(skeleton: np.ndarray, width: float, height: float) -> np.ndarray:
    """Correct normalized image coordinates before geometric computation.

    Input/output: (F, 17, C). Formula: x' = x * (width / height), y' = y.
    This makes horizontal and vertical coordinate units geometrically comparable.
    """
    arr = as_skeleton_array(skeleton)
    if height <= 0:
        raise ValueError("height must be positive")
    out = arr.copy()
    out[..., 0] *= float(width) / float(height)
    return out


def center_on_hips(skeleton: np.ndarray) -> np.ndarray:
    """Subtract the per-frame hip midpoint from every joint.

    Input: (F, 17, C). Output: (F, 17, 2). Formula:
    centered[f, j] = xy[f, j] - 0.5 * (xy[f, left_hip] + xy[f, right_hip]).
    """
    coords = xy(skeleton)
    hip_midpoint = 0.5 * (coords[:, LEFT_HIP, :] + coords[:, RIGHT_HIP, :])
    return coords - hip_midpoint[:, None, :]


def normalize_scale(skeleton: np.ndarray, eps: float = 1e-6) -> np.ndarray:
    """Divide coordinates by torso length per frame.

    Input/output: (F, 17, 2). Formula: xy' = xy / max(||nose - hip_midpoint||, eps).
    If the nose distance degenerates, shoulder-to-hip distance is used as fallback.
    """
    coords = xy(skeleton)
    hip_midpoint = 0.5 * (coords[:, LEFT_HIP, :] + coords[:, RIGHT_HIP, :])
    shoulder_midpoint = 0.5 * (coords[:, LEFT_SHOULDER, :] + coords[:, RIGHT_SHOULDER, :])
    torso = np.linalg.norm(coords[:, NOSE, :] - hip_midpoint, axis=1)
    fallback = np.linalg.norm(shoulder_midpoint - hip_midpoint, axis=1)
    scale = np.where(torso > eps, torso, fallback)
    scale = np.maximum(scale, eps)
    return coords / scale[:, None, None]


def compute_joint_angles(
    skeleton: np.ndarray, triplets: tuple[tuple[int, int, int], ...] = JOINT_ANGLE_TRIPLETS
) -> np.ndarray:
    """Compute cosine joint angles for anatomical triplets.

    Input: (F, 17, 2). Output: (F, len(triplets)).
    Formula for (a, b, c): cos(theta) = dot(a-b, c-b) / (||a-b|| * ||c-b||).
    The middle joint b is the angle vertex; values are clipped to [-1, 1].
    """
    coords = xy(skeleton)
    angles = []
    for a, b, c in triplets:
        ba = coords[:, a, :] - coords[:, b, :]
        bc = coords[:, c, :] - coords[:, b, :]
        denom = np.linalg.norm(ba, axis=1) * np.linalg.norm(bc, axis=1)
        cosine = np.divide(
            np.sum(ba * bc, axis=1),
            denom,
            out=np.zeros(coords.shape[0], dtype=np.float64),
            where=denom > 1e-8,
        )
        angles.append(np.clip(cosine, -1.0, 1.0))
    return np.stack(angles, axis=1)


def compute_bone_vectors(
    skeleton: np.ndarray, bone_pairs: tuple[tuple[int, int], ...] = BONE_PAIRS
) -> np.ndarray:
    """Compute parent-to-child displacement vectors.

    Input: (F, 17, 2). Output: (F, len(bone_pairs) * 2).
    Formula for (parent, child): vector = xy[child] - xy[parent].
    """
    coords = xy(skeleton)
    vectors = [coords[:, child, :] - coords[:, parent, :] for parent, child in bone_pairs]
    return np.concatenate(vectors, axis=1)


def compute_velocities(skeleton: np.ndarray, fps: float) -> np.ndarray:
    """Compute first temporal derivatives with central differences.

    Input: (F, J, 2). Output: (F, J, 2). Formula approximates dx/dt via np.gradient(x, 1/fps).
    """
    coords = xy(skeleton)
    return np.gradient(coords, 1.0 / float(fps), axis=0, edge_order=1)


def compute_accelerations(skeleton: np.ndarray, fps: float) -> np.ndarray:
    """Compute second temporal derivatives of keypoint coordinates.

    Input: (F, J, 2). Output: (F, J, 2). Formula: d2x/dt2 = gradient(gradient(x)).
    """
    return np.gradient(compute_velocities(skeleton, fps), 1.0 / float(fps), axis=0, edge_order=1)


def compute_angular_velocities(angles: np.ndarray, fps: float) -> np.ndarray:
    """Compute first temporal derivative of angle-cosine features.

    Input: (F, A). Output: (F, A). Formula approximates d(cos(theta))/dt by central differences.
    """
    return np.gradient(np.asarray(angles, dtype=np.float64), 1.0 / float(fps), axis=0, edge_order=1)


def build_feature_tensor(skeleton: np.ndarray, fps: float) -> np.ndarray:
    """Build the public 94-dimensional feature tensor.

    Input: cleaned, hip-centered, scale-normalized skeleton with shape (F, 17, 2).
    Output: (F, 94) with [34 coords, 12 angle cosines, 20 bone-vector values,
    16 selected keypoint velocity values, 12 angular velocity values].
    Correctness note: this function assumes aspect-ratio correction and biomechanical
    normalization have already been applied.
    """
    coords = xy(skeleton)
    angles = compute_joint_angles(coords)
    bone_vectors = compute_bone_vectors(coords)
    velocities = compute_velocities(coords, fps)[:, KEY_JOINTS, :].reshape(coords.shape[0], -1)
    angular_velocities = compute_angular_velocities(angles, fps)
    return np.concatenate(
        [
            coords.reshape(coords.shape[0], -1),
            angles,
            bone_vectors,
            velocities,
            angular_velocities,
        ],
        axis=1,
    )


def resample_to_length(features: np.ndarray, target_length: int) -> np.ndarray:
    """Linearly resample a sequence to a fixed frame count.

    Input: (F, D). Output: (target_length, D). The original sequence is treated as samples
    over normalized time [0, 1], and every feature dimension is interpolated independently.
    """
    arr = np.asarray(features, dtype=np.float64)
    if arr.ndim != 2:
        raise ValueError(f"Expected features with shape (F, D), got {arr.shape}")
    if target_length <= 0:
        raise ValueError("target_length must be positive")
    if arr.shape[0] == target_length:
        return arr.copy()
    if arr.shape[0] == 0:
        raise ValueError("Cannot resample an empty sequence")
    if arr.shape[0] == 1:
        return np.repeat(arr, target_length, axis=0)

    old_t = np.linspace(0.0, 1.0, arr.shape[0])
    new_t = np.linspace(0.0, 1.0, target_length)
    out = np.empty((target_length, arr.shape[1]), dtype=np.float64)
    for dim in range(arr.shape[1]):
        out[:, dim] = np.interp(new_t, old_t, arr[:, dim])
    return out