Biomechanical Indices from Pose Angles

# 3) Imports & Environment Check
import sys
from pathlib import Path
import numpy as np
import pandas as pd

print("Python:", sys.version)
print("Pandas:", pd.__version__)
print("NumPy :", np.__version__)

# Set your angles CSV path (from Notebook 2)
angles_csv_path = r""  # put the path of your csv file in between the two double quotes. e.g., "/path/to/outputs/video_pose2d_angles.csv"

if not angles_csv_path or not str(angles_csv_path).strip():
    raise ValueError("Please set `angles_csv_path` to a valid file.")

angles_csv_path = Path(angles_csv_path).expanduser().resolve()
if not angles_csv_path.exists() or not angles_csv_path.is_file():
    raise FileNotFoundError(f"Angles CSV not found: {angles_csv_path}")

angles_df = pd.read_csv(angles_csv_path)
print("Loaded:", angles_csv_path, "shape:", angles_df.shape)
print("Angle columns:", [c for c in angles_df.columns if c.startswith("angle_")][:10], "…")
print("Confidence columns:", [c for c in angles_df.columns if c.startswith("confidence_")][:10], "…")
angles_df.head(3)

# Selection / Reduction Parameters — YOU CAN EDIT THIS BLOCK

# --- Choose ONE mode: "all_frames" | "angle_thresholded" | "running_avg"
selection_mode = "all_frames"

# --- Angles to consider (must exist as columns in angles_df, e.g., "angle_left_knee")
selected_angles = [
    "angle_right_body",
    "angle_right_elbow",
    "angle_right_knee",
    "angle_right_hip",
]

# --- Confidence gate (optional): keep frames where each angle's confidence is in this set
allowed_conf_levels = {"good"}   # e.g., {"good"}, or {"good","low"}

# --- If your CSV contains multiple videos, you can focus on one (or leave None)
selected_video = None            # e.g., "myvideo.mp4" or None to keep all

# --- Angle-thresholded mode: define per-angle conditions
# Supported operators: ">", ">=", "<", "<=", "==", "!="
# threshold_logic: "all" (all conditions must hold) or "any" (at least one)
angle_thresholds = {
    # "angle_left_knee": (">=", 140),
    # "angle_right_knee": (">=", 140),
}
threshold_logic = "all"          # "all" | "any"

# --- Running-average mode: rolling window size (frames) and centering
running_window = 11              # odd number recommended (e.g., 11 means ±5 frames)
running_center = True            # center the window
running_min_periods = 1          # minimum frames to compute a mean

import operator
import pandas as pd
import numpy as np

# Defensive checks
missing_cols = [c for c in selected_angles if c not in angles_df.columns]
if missing_cols:
    raise ValueError(f"Selected angle columns not found: {missing_cols}")

# 1) subset by video (optional)
work_df = angles_df.copy()
if selected_video is not None:
    work_df = work_df[work_df["video"] == selected_video].copy()

# 2) sort by frame (important for running averages)
if "frame" not in work_df.columns:
    raise ValueError("Input CSV must have a 'frame' column.")
work_df = work_df.sort_values(["video","frame"] if "video" in work_df.columns else ["frame"])

# 3) optional confidence filter per angle
def _apply_conf_filter_per_angle(df, angle_col, allowed):
    conf_col = "confidence_" + angle_col.replace("angle_", "", 1)
    if conf_col in df.columns:
        return df[df[conf_col].isin(allowed)]
    return df

if allowed_conf_levels:
    keep_idx = pd.Series(True, index=work_df.index)
    for ang in selected_angles:
        filtered = _apply_conf_filter_per_angle(work_df, ang, allowed_conf_levels)
        keep_idx &= work_df.index.isin(filtered.index)
    work_df = work_df.loc[keep_idx].copy()

# 4) selection modes
if selection_mode == "all_frames":
    # keep angles as-is
    pass

elif selection_mode == "angle_thresholded":
    # Build a boolean mask from angle_thresholds using AND/OR logic
    if not angle_thresholds:
        raise ValueError("You selected 'angle_thresholded' but did not define any angle_thresholds.")
    ops = {">": operator.gt, ">=": operator.ge, "<": operator.lt, "<=": operator.le, "==": operator.eq, "!=": operator.ne}

    masks = []
    for ang, (op_str, val) in angle_thresholds.items():
        if ang not in work_df.columns:
            raise ValueError(f"Angle '{ang}' not found in DataFrame.")
        if op_str not in ops:
            raise ValueError(f"Unsupported operator '{op_str}' for '{ang}'.")
        masks.append(ops[op_str](work_df[ang].astype(float), float(val)))

    if threshold_logic == "all":
        keep = np.logical_and.reduce(masks)
    elif threshold_logic == "any":
        keep = np.logical_or.reduce(masks)
    else:
        raise ValueError("threshold_logic must be 'all' or 'any'.")

    work_df = work_df[keep].copy()

elif selection_mode == "running_avg":
    # Replace each selected angle with its rolling mean over the whole dataset
    # Do this video-by-video if a 'video' column exists
    if running_window is None or running_window < 1:
        raise ValueError("running_window must be a positive integer.")
    if running_min_periods is None:
        running_min_periods = 1

    if "video" in work_df.columns:
        work_df = (
            work_df
            .groupby("video", group_keys=False)
            .apply(lambda d: d.assign(**{
                ang: d[ang].astype(float).rolling(
                    window=running_window, center=running_center, min_periods=running_min_periods
                ).mean()
                for ang in selected_angles
            }))
        )
    else:
        for ang in selected_angles:
            work_df[ang] = work_df[ang].astype(float).rolling(
                window=running_window, center=running_center, min_periods=running_min_periods
            ).mean()
else:
    raise ValueError("selection_mode must be one of: 'all_frames', 'angle_thresholded', 'running_avg'")

print(f"Selection mode: {selection_mode}")
print(f"Rows after selection: {len(work_df)}")
display(work_df.head(8)[['video','frame'] + selected_angles if 'video' in work_df.columns else ['frame'] + selected_angles])

# Reuse your discrete_rules and continuous_configs from earlier cells
# If you haven't defined them yet, do that first (see the dedicated sections).

results = work_df[['video','frame']].copy() if 'video' in work_df.columns else work_df[['frame']].copy()
for ang in selected_angles:
    if ang in work_df.columns:
        results[ang] = work_df[ang].astype(float)

# Apply discrete labels where rules exist
if 'discrete_rules' in globals():
    for col, rules in discrete_rules.items():
        if col in results.columns:
            # apply_discrete_rules defined earlier in the notebook
            results[f"label_{col}"] = apply_discrete_rules(results[col], rules)

# Apply continuous scores where configs exist
if 'continuous_configs' in globals():
    # apply_continuous_scores defined earlier in the notebook
    results = apply_continuous_scores(results, continuous_configs)

print("Computed indices on the selected angle set:")
display(results.head(10))

# Reuse your discrete_rules and continuous_configs from earlier cells
# If you haven't defined them yet, do that first (see the dedicated sections).

results = work_df[['video','frame']].copy() if 'video' in work_df.columns else work_df[['frame']].copy()
for ang in selected_angles:
    if ang in work_df.columns:
        results[ang] = work_df[ang].astype(float)

# Apply discrete labels where rules exist
if 'discrete_rules' in globals():
    for col, rules in discrete_rules.items():
        if col in results.columns:
            # apply_discrete_rules defined earlier in the notebook
            results[f"label_{col}"] = apply_discrete_rules(results[col], rules)

# Apply continuous scores where configs exist
if 'continuous_configs' in globals():
    # apply_continuous_scores defined earlier in the notebook
    results = apply_continuous_scores(results, continuous_configs)

print("Computed indices on the selected angle set:")
display(results.head(10))

out_dir = angles_csv_path.parent
suffix = {
    "all_frames": "all",
    "angle_thresholded": "thresh",
    "running_avg": f"runavg_w{running_window}"
}[selection_mode]

out_csv = out_dir / f"{angles_csv_path.stem}_indices_{suffix}.csv"
results.to_csv(out_csv, index=False)
print("Saved indices to:", out_csv)