formscout/analysis/charts.py

"""
Matplotlib chart generators for the screening report.

Every function returns a PNG path on success or None on failure (never raises),
so a chart problem degrades the report but never blocks scoring. Charts use the
Silas palette and an Agg backend so they render headless on the Space.
"""
from __future__ import annotations

import logging

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402

from formscout.analysis.relevant_joints import COCO_NAMES  # noqa: E402

logger = logging.getLogger(__name__)

TEAL = "#2b8a8a"
GOLD = "#e0a43b"
SAGE = "#9cbcad"
INK = "#243a34"
RED = "#d9534f"
_PALETTE = [TEAL, GOLD, SAGE, "#7a5ca0", "#c2683c", "#3c8dbc"]


def _save(fig, out_png: str) -> str | None:
    try:
        fig.savefig(out_png, dpi=110, bbox_inches="tight", facecolor="white")
        return out_png
    except Exception as e:
        logger.warning("chart save failed: %s", e)
        return None
    finally:
        plt.close(fig)


def angle_over_time(series: dict, primary: str | None, governing_idx: int | None,
                    out_png: str, title: str = "Joint angle over time") -> str | None:
    """Angle-vs-frame for the relevant angles; primary emphasised, key-frame marked."""
    try:
        if not series:
            return None
        fig, ax = plt.subplots(figsize=(6.4, 3.2))
        for i, (name, vals) in enumerate(series.items()):
            arr = np.array(vals, dtype=float)
            is_primary = name == primary
            ax.plot(np.arange(len(arr)), arr,
                    color=(TEAL if is_primary else _PALETTE[i % len(_PALETTE)]),
                    lw=2.4 if is_primary else 1.3,
                    alpha=1.0 if is_primary else 0.6,
                    label=name.replace("_", " ") + (" ★" if is_primary else ""))
        if governing_idx is not None:
            ax.axvline(governing_idx, color=GOLD, ls="--", lw=1.5, label="key frame")
        ax.set_xlabel("frame")
        ax.set_ylabel("degrees")
        ax.set_title(title, color=INK)
        ax.legend(fontsize=7, loc="best")
        ax.grid(True, alpha=0.2)
        return _save(fig, out_png)
    except Exception as e:
        logger.warning("angle_over_time failed: %s", e)
        return None


def velocity_profile(keypoints: list, fps: float, joints: list[int],
                     out_png: str, title: str = "Joint speed over time") -> str | None:
    """Per-frame speed (px/s) of the relevant joints across the clip."""
    try:
        from formscout.agents.visualizer import compute_joint_velocity
        vel = compute_joint_velocity(keypoints, fps or 30.0)
        plot_joints = [j for j in joints if j in vel] or list(vel.keys())[:4]
        if not plot_joints:
            return None
        fig, ax = plt.subplots(figsize=(6.4, 3.2))
        for i, j in enumerate(plot_joints):
            ax.plot(vel[j], color=_PALETTE[i % len(_PALETTE)], lw=1.6,
                    label=COCO_NAMES.get(j, str(j)).replace("_", " "))
        ax.set_xlabel("frame")
        ax.set_ylabel("speed (px/s)")
        ax.set_title(title, color=INK)
        ax.legend(fontsize=7, loc="best")
        ax.grid(True, alpha=0.2)
        return _save(fig, out_png)
    except Exception as e:
        logger.warning("velocity_profile failed: %s", e)
        return None


def laban_radar(effort: dict, out_png: str, title: str = "Laban Effort") -> str | None:
    """4-axis radar of the Effort factors (Space, Weight, Time, Flow)."""
    try:
        axes_order = ["space", "weight", "time", "flow"]
        labels = ["Space\n(direct)", "Weight\n(strong)", "Time\n(sudden)", "Flow\n(free)"]
        vals = [float(effort.get(k, 0.0)) for k in axes_order]
        angles = np.linspace(0, 2 * np.pi, len(axes_order), endpoint=False).tolist()
        vals_loop = vals + vals[:1]
        angles_loop = angles + angles[:1]

        fig, ax = plt.subplots(figsize=(4.2, 4.2), subplot_kw={"polar": True})
        ax.plot(angles_loop, vals_loop, color=TEAL, lw=2)
        ax.fill(angles_loop, vals_loop, color=TEAL, alpha=0.25)
        ax.set_xticks(angles)
        ax.set_xticklabels(labels, fontsize=8, color=INK)
        ax.set_ylim(0, 1)
        ax.set_yticks([0.25, 0.5, 0.75, 1.0])
        ax.set_yticklabels(["", "0.5", "", "1.0"], fontsize=7)
        ax.set_title(title, color=INK, pad=18)
        return _save(fig, out_png)
    except Exception as e:
        logger.warning("laban_radar failed: %s", e)
        return None


def flexion_bars(flexion: dict, out_png: str,
                 title: str = "Relevant joint flexion") -> str | None:
    """Horizontal bars of relevant joint angles (deg) at the key frame."""
    try:
        if not flexion:
            return None
        names = [n.replace("_", " ") for n in flexion]
        degs = [flexion[n]["deg"] for n in flexion]
        colors = [TEAL if d >= 160 else GOLD if d >= 110 else RED for d in degs]
        fig, ax = plt.subplots(figsize=(6.0, max(1.6, 0.5 * len(names) + 0.8)))
        y = np.arange(len(names))
        ax.barh(y, degs, color=colors)
        ax.set_yticks(y)
        ax.set_yticklabels(names, fontsize=8)
        ax.set_xlim(0, 200)
        ax.axvline(160, color=SAGE, ls=":", lw=1)
        for yi, d in zip(y, degs):
            ax.text(d + 3, yi, f"{d:.0f}°", va="center", fontsize=8, color=INK)
        ax.set_xlabel("interior angle (°)  ·  higher = more open")
        ax.set_title(title, color=INK)
        ax.invert_yaxis()
        return _save(fig, out_png)
    except Exception as e:
        logger.warning("flexion_bars failed: %s", e)
        return None


def symmetry_bars(asymmetries: list, out_png: str,
                  title: str = "Left / right symmetry") -> str | None:
    """Grouped L vs R score bars for bilateral tests."""
    try:
        rows = [a for a in asymmetries
                if a.get("left_score") is not None and a.get("right_score") is not None]
        if not rows:
            return None
        names = [a["test"].replace("_", " ") for a in rows]
        left = [a["left_score"] for a in rows]
        right = [a["right_score"] for a in rows]
        x = np.arange(len(names))
        w = 0.36
        fig, ax = plt.subplots(figsize=(6.4, 3.2))
        ax.bar(x - w / 2, left, w, color=TEAL, label="left")
        ax.bar(x + w / 2, right, w, color=GOLD, label="right")
        ax.set_xticks(x)
        ax.set_xticklabels(names, fontsize=8, rotation=15, ha="right")
        ax.set_ylim(0, 3.4)
        ax.set_ylabel("score (0–3)")
        ax.set_title(title, color=INK)
        ax.legend(fontsize=8)
        ax.grid(True, axis="y", alpha=0.2)
        return _save(fig, out_png)
    except Exception as e:
        logger.warning("symmetry_bars failed: %s", e)
        return None