plot_knee_angles.py

"""
Knee IMU Angle Plotter
======================
Plots pitch (primary flexion/extension), roll (varus/valgus),
and yaw (internal/external rotation) angles from two IMUs
placed on the thigh and shin, vs time.

Usage:
    python plot_knee_angles.py                        # uses default CSV path
    python plot_knee_angles.py my_data.csv            # custom CSV path
    python plot_knee_angles.py my_data.csv output.png # custom CSV + save figure
"""

import sys
import os
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.ticker import AutoMinorLocator
import numpy as np

# ── Config ─────────────────────────────────────────────────────────────────────
DEFAULT_CSV = "sensor-data-2026-03-04T15-17-26.csv"
SMOOTHING_WINDOW = 15          # rolling-average window (samples); set to 1 to disable
FIGURE_SIZE      = (14, 9)
PRIMARY_COLOR    = "#2563EB"   # blue  – flexion (pitch)
SECONDARY_COLOR  = "#16A34A"   # green – roll
TERTIARY_COLOR   = "#DC2626"   # red   – yaw
ALPHA_RAW        = 0.25        # opacity of raw data
ALPHA_SMOOTH     = 1.0         # opacity of smoothed line
# ───────────────────────────────────────────────────────────────────────────────


def load_data(csv_path: str) -> pd.DataFrame:
    df = pd.read_csv(csv_path)

    # Normalise column names to lowercase
    df.columns = [c.strip().lower() for c in df.columns]

    required = {"timestamp", "pitch", "roll", "yaw"}
    missing = required - set(df.columns)
    if missing:
        raise ValueError(f"CSV is missing expected columns: {missing}\n"
                         f"Found columns: {list(df.columns)}")

    # Convert timestamp (ms) to elapsed seconds
    df["time_s"] = (df["timestamp"] - df["timestamp"].iloc[0]) / 1000.0

    # Smooth each angle channel
    for col in ("pitch", "roll", "yaw"):
        df[f"{col}_smooth"] = (
            df[col]
            .rolling(SMOOTHING_WINDOW, center=True, min_periods=1)
            .mean()
        )

    return df


def compute_stats(series: pd.Series) -> dict:
    return {
        "min":   series.min(),
        "max":   series.max(),
        "range": series.max() - series.min(),
        "mean":  series.mean(),
    }


def plot(df: pd.DataFrame, save_path: str | None = None) -> None:
    total_time = df["time_s"].iloc[-1]

    fig = plt.figure(figsize=FIGURE_SIZE, facecolor="#F8FAFC")
    fig.suptitle(
        "Knee IMU — Relative Joint Angles vs Time",
        fontsize=16, fontweight="bold", color="#1E293B", y=0.98
    )

    gs = gridspec.GridSpec(
        3, 1, figure=fig,
        hspace=0.08,         # tight spacing between subplots
        top=0.92, bottom=0.08, left=0.08, right=0.97
    )

    axes = [fig.add_subplot(gs[i]) for i in range(3)]

    channels = [
        {
            "col":    "pitch",
            "label":  "Pitch — Flexion / Extension  (primary)",
            "color":  PRIMARY_COLOR,
            "ylabel": "Pitch (°)",
            "annot":  "Main squat angle",
        },
        {
            "col":    "roll",
            "label":  "Roll — Varus / Valgus  (secondary)",
            "color":  SECONDARY_COLOR,
            "ylabel": "Roll (°)",
            "annot":  "Side-to-side tilt",
        },
        {
            "col":    "yaw",
            "label":  "Yaw — Internal / External Rotation  (secondary)",
            "color":  TERTIARY_COLOR,
            "ylabel": "Yaw (°)",
            "annot":  "Axial rotation",
        },
    ]

    for ax, ch in zip(axes, channels):
        raw    = df[ch["col"]]
        smooth = df[f"{ch['col']}_smooth"]
        t      = df["time_s"]

        # Raw data (faint fill + scatter)
        ax.fill_between(t, raw, alpha=0.08, color=ch["color"])
        ax.plot(t, raw, color=ch["color"], alpha=ALPHA_RAW,
                linewidth=0.8, label="Raw")

        # Smoothed line
        ax.plot(t, smooth, color=ch["color"], alpha=ALPHA_SMOOTH,
                linewidth=2.0, label=f"Smoothed (n={SMOOTHING_WINDOW})")

        # Zero reference line
        ax.axhline(0, color="#94A3B8", linewidth=0.6, linestyle="--", zorder=0)

        # Stats annotation
        stats = compute_stats(raw)
        stat_text = (
            f"min {stats['min']:.1f}°   "
            f"max {stats['max']:.1f}°   "
            f"range {stats['range']:.1f}°   "
            f"mean {stats['mean']:.1f}°"
        )
        ax.text(0.01, 0.97, stat_text,
                transform=ax.transAxes, fontsize=7.5,
                verticalalignment="top", color="#475569",
                fontfamily="monospace")

        # Channel label
        ax.set_ylabel(ch["ylabel"], fontsize=10, color="#1E293B", labelpad=6)
        ax.set_title(ch["label"], fontsize=10.5, color=ch["color"],
                     fontweight="semibold", loc="left", pad=3)

        # Grid & ticks
        ax.set_facecolor("#F1F5F9")
        ax.grid(True, which="major", color="#CBD5E1", linewidth=0.6)
        ax.grid(True, which="minor", color="#E2E8F0", linewidth=0.3)
        ax.yaxis.set_minor_locator(AutoMinorLocator())
        ax.xaxis.set_minor_locator(AutoMinorLocator())
        ax.tick_params(labelsize=8, colors="#475569")
        for spine in ax.spines.values():
            spine.set_edgecolor("#CBD5E1")

        ax.set_xlim(0, total_time)
        ax.legend(fontsize=8, loc="upper right",
                  framealpha=0.7, edgecolor="#CBD5E1")

        # Only bottom axis gets x-label
        if ax is axes[-1]:
            ax.set_xlabel("Time (s)", fontsize=10, color="#1E293B")
        else:
            ax.set_xticklabels([])

    # Caption
    fig.text(
        0.5, 0.01,
        f"Duration: {total_time:.1f} s   |   Samples: {len(df)}   |   "
        f"Sample rate ≈ {len(df)/total_time:.0f} Hz",
        ha="center", fontsize=8, color="#94A3B8"
    )

    if save_path:
        fig.savefig(save_path, dpi=150, bbox_inches="tight")
        print(f"Figure saved → {save_path}")
    else:
        plt.show()


# ── Entry point ────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    csv_path  = sys.argv[1] if len(sys.argv) > 1 else DEFAULT_CSV
    save_path = sys.argv[2] if len(sys.argv) > 2 else None

    if not os.path.isfile(csv_path):
        print(f"ERROR: File not found — {csv_path}")
        sys.exit(1)

    print(f"Loading data from: {csv_path}")
    df = load_data(csv_path)
    print(f"  Loaded {len(df)} samples  |  "
          f"Duration: {df['time_s'].iloc[-1]:.1f} s")

    plot(df, save_path)