app.py

"""
FormScout — Gradio app entrypoint.
Screening aid for Functional Movement Screen (FMS) scoring.
NOT a diagnosis. NOT an injury predictor.

Custom scout/trail themed UI with score dial, pipeline visualization,
rubric breakdown, and persistent safety banner.
"""
from __future__ import annotations

import os
import tempfile

import gradio as gr

from formscout.pipeline import Director
from formscout.rubric import score_test
from formscout.ui.theme import formscout_theme, FORMSCOUT_CSS
from formscout import config
from formscout import session as session_mod
from formscout.startup import ensure_checkpoints

# ─── ZeroGPU ──────────────────────────────────────────────────────────────────
# On an HF Spaces ZeroGPU runtime the heavy analysis MUST run inside an
# @spaces.GPU function, and that function must already exist at import time:
# ZeroGPU scans for one during startup and aborts the Space with
# "No @spaces.GPU function detected during startup" if none is registered.
# We decorate process_video (the Start Analysis handler) so a single GPU window
# covers the whole pipeline — pose, optional 3D, and the Qwen3-VL judge. Off a
# ZeroGPU Space the `spaces` package is absent (or its decorator is effect-free),
# so local runs and CPU Spaces are unaffected.
try:
    import spaces

    gpu_task = spaces.GPU(duration=config.ZEROGPU_DURATION)
except Exception:  # local dev / non-ZeroGPU — decorate as a no-op
    def gpu_task(fn):
        return fn

ensure_checkpoints()


# ─── Constants ───────────────────────────────────────────────────────────────

DISCLAIMER = (
    "⚠️ **Screening aid — not a diagnosis. "
    "Pain or clearing tests require a clinician.**"
)

FMS_TESTS = [
    ("Deep Squat", "deep_squat"),
    ("Hurdle Step", "hurdle_step"),
    ("In-Line Lunge", "inline_lunge"),
    ("Shoulder Mobility", "shoulder_mobility"),
    ("Active Straight-Leg Raise", "active_slr"),
    ("Trunk Stability Push-Up", "trunk_stability_pushup"),
    ("Rotary Stability", "rotary_stability"),
]

SCORE_DESCRIPTIONS = {
    3: "Movement performed to criterion — no compensation",
    2: "Movement completed with compensation or regression",
    1: "Unable to perform the movement pattern",
    0: "Pain reported — clinician referral required",
}

# Per-test illustration shown as a movement guide (test_key -> filename).
_ASSET_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "assets", "screenings")
SCREENING_ASSETS = {
    "deep_squat": "deep-squat.png",
    "hurdle_step": "hurdle-step.png",
    "inline_lunge": "inline-lunge.png",
    "shoulder_mobility": "shoulder-mobility.png",
    "active_slr": "active-straight-leg-raise.png",
    "trunk_stability_pushup": "trunk-stability-push-up.png",
    "rotary_stability": "rotary-stability.png",
}


def _screening_image(test_display_name: str) -> str | None:
    """Resolve the movement-guide image path for a dropdown display name."""
    key = {name: val for name, val in FMS_TESTS}.get(test_display_name)
    fname = SCREENING_ASSETS.get(key)
    if not fname:
        return None
    path = os.path.join(_ASSET_DIR, fname)
    return path if os.path.exists(path) else None


# ─── Processing ──────────────────────────────────────────────────────────────

@gpu_task
def process_video(video_path: str, test_name: str, side: str, model_key: str,
                  layers: list[str], session_state):
    """Analyse one clip and accumulate it into the screening session.

    Decorated with @spaces.GPU on ZeroGPU: the whole pipeline (pose, optional 3D,
    Qwen3-VL judge) runs inside one GPU window. The decorator is a no-op off-Space.
    """
    if not video_path:
        return (
            session_state, _render_empty_state(), "Upload a video to begin analysis.",
            "", "", None, "", _render_session_table(session_state),
            gr.update(visible=False), gr.update(visible=False),
        )

    if session_state is None:
        session_state = session_mod.new_session()

    director = Director()
    state = director.run(video_path, test_name=test_name, side=side, model_key=model_key)

    score_html = _render_empty_state()
    score_details = ""

    if state.features:
        result = score_test(state.features)
        judge = state.judge
        if judge and judge.score is not None:
            score_html = _render_score_card(judge.score, judge.confidence, judge.needs_human)
            score_details = _render_score_details_judge(judge, result, state.features)
        elif judge and judge.needs_human:
            score_html = _render_score_card(0, 0, True)
            score_details = f"### Needs Clinician Review\n{judge.rationale}"
        else:
            score_html = _render_score_card(result.score, result.confidence, result.needs_human)
            score_details = _render_score_details(result, state.features)

        # Accumulate into the session (only when we have a real analysis)
        if state.ingest and state.pose2d and state.judge:
            draw_trails = "trails" in {lbl.lower().replace(" ", "_") for lbl in (layers or [])}
            try:
                session_mod.add_analysis(
                    session_state, ingest=state.ingest, pose2d=state.pose2d,
                    features=state.features, judge=state.judge,
                    test_name=test_name, side=side, draw_trails=draw_trails,
                )
            except Exception as e:
                state.warnings.append(f"session accumulation failed: {e}")

    pipeline_md = _render_pipeline_status(state)
    alerts = _render_alerts(state)

    overlay_path = None
    vel_summary = ""
    layer_set = {lbl.lower().replace(" ", "_") for lbl in (layers or [])}
    if layer_set and state.ingest and state.pose2d:
        try:
            from formscout.agents.visualizer import PoseVisualizer, build_velocity_summary
            vis = PoseVisualizer()
            with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
                out_path = f.name
            overlay_path = vis.render_video(state.ingest, state.pose2d, layer_set, out_path)
            if overlay_path:
                vel_summary = build_velocity_summary(state.pose2d.keypoints, vis.last_velocities)
        except Exception as e:
            alerts = (alerts or "") + f"\n⚠️ Visualizer error: {e}"

    has_entries = bool(session_state and session_state.entries)
    return (
        session_state, score_html, pipeline_md, score_details, alerts,
        overlay_path, vel_summary, _render_session_table(session_state),
        gr.update(visible=has_entries), gr.update(visible=has_entries),
    )


def _render_score_card(score: int, confidence: float, needs_human: bool) -> str:
    """Render the score dial as HTML."""
    if needs_human:
        return """
        <div class="score-card needs-review">
            <div style="font-size: 1.2em; color: #cf922a; margin-bottom: 8px;">⚠️ Needs Clinician Review</div>
            <div style="font-size: 0.9em; color: #4a5f57;">Pain or clearing test detected — cannot auto-score</div>
        </div>
        """

    conf_pct = int(confidence * 100)
    conf_color = "#2b8a8a" if confidence >= 0.7 else "#cf922a" if confidence >= 0.4 else "#d9534f"

    return f"""
    <div class="score-card">
        <div class="score-value">{score}/3</div>
        <div style="font-size: 0.95em; color: #4a5f57; margin-top: 4px;">
            {SCORE_DESCRIPTIONS.get(score, '')}
        </div>
        <div style="margin-top: 12px;">
            <div style="display: flex; justify-content: space-between; font-size: 0.8em; color: #6b7d75;">
                <span>Confidence</span>
                <span style="color: {conf_color};">{conf_pct}%</span>
            </div>
            <div class="confidence-bar">
                <div class="confidence-fill" style="width: {conf_pct}%;"></div>
            </div>
        </div>
    </div>
    """


def _render_empty_state() -> str:
    """Render placeholder when no video processed yet."""
    return """
    <div class="score-card" style="opacity: 0.6;">
        <div style="font-size: 2em; margin-bottom: 8px;">🏔️</div>
        <div style="color: #6b7d75;">Upload a video to begin</div>
    </div>
    """


def _render_score_details(result, features) -> str:
    """Render the rubric breakdown."""
    parts = [f"### Rationale\n{result.rationale}\n"]

    if features.angles:
        parts.append("### Measurements")
        for key, val in features.angles.items():
            label = key.replace("_", " ").title()
            parts.append(f"- **{label}:** {val:.1f}°")

    if features.alignments:
        parts.append("\n### Alignment Checks")
        for key, val in features.alignments.items():
            label = key.replace("_", " ").title()
            icon = "✓" if val else "✗"
            parts.append(f"- {icon} {label}")

    if features.view == "2d":
        parts.append(
            "\n> ⚠️ *2D estimate — angles are camera-angle dependent. "
            "For best accuracy, film from the side at hip height.*"
        )

    return "\n".join(parts)


def _render_score_details_judge(judge, rubric, features) -> str:
    """Render judge + rubric combined breakdown."""
    parts = [f"### Judge Rationale\n{judge.rationale}\n"]

    if judge.compensation_tags:
        parts.append(f"**Compensations:** {', '.join(judge.compensation_tags)}")
    if judge.corrective_hint:
        parts.append(f"**Corrective:** {judge.corrective_hint}")

    parts.append(f"\n### Rubric Score: {rubric.score}/3")
    parts.append(f"*{rubric.rationale}*")

    if features.angles:
        parts.append("\n### Measurements")
        for key, val in features.angles.items():
            label = key.replace("_", " ").title()
            parts.append(f"- **{label}:** {val:.1f}°" if isinstance(val, float) else f"- **{label}:** {val}")

    if features.symmetry_delta is not None:
        parts.append(f"\n### Asymmetry\n- **L/R Delta:** {features.symmetry_delta:.1f}°")

    if features.view == "2d":
        parts.append(
            "\n> ⚠️ *2D estimate — angles are camera-angle dependent.*"
        )

    return "\n".join(parts)


def _render_pipeline_status(state) -> str:
    """Render pipeline step summary."""
    parts = []
    if state.ingest:
        parts.append(
            f"📹 **Ingest:** {len(state.ingest.frames)} frames · "
            f"{state.ingest.fps:.0f}fps · {state.ingest.duration:.1f}s · "
            f"{state.ingest.width}×{state.ingest.height}"
        )
    if state.pose2d:
        n = sum(1 for kps in state.pose2d.keypoints if kps)
        parts.append(
            f"🦴 **Pose2D:** {n}/{len(state.pose2d.keypoints)} frames detected · "
            f"conf={state.pose2d.confidence:.0%}"
        )
    if state.body3d:
        if state.body3d.used:
            parts.append(f"🧊 **Body3D:** active · conf={state.body3d.confidence:.0%}")
        else:
            parts.append("🧊 **Body3D:** 2D-only path (normal)")
    if state.features:
        parts.append(
            f"📐 **Biomechanics:** view={state.features.view} · "
            f"conf={state.features.confidence:.0%}"
        )
    return "\n\n".join(parts) if parts else "*Processing...*"


def _render_alerts(state) -> str:
    """Render errors and warnings."""
    parts = []
    if state.errors:
        for e in state.errors:
            parts.append(f"🚨 {e}")
    if state.warnings:
        for w in state.warnings:
            parts.append(f"⚠️ {w}")
    return "\n\n".join(parts)


def _render_session_table(session_state) -> str:
    """Render the accumulated 'Session so far' table as markdown."""
    if not session_state or not session_state.entries:
        return "*No clips analysed yet.*"
    lines = ["| Test | Side | Score | Status |", "|---|---|---|---|"]
    for e in session_state.entries:
        test = e.test_name.replace("_", " ").title()
        side = e.side if e.side in ("left", "right") else "—"
        if e.needs_human:
            score, status = "—", "⚠️ Clinician review"
        else:
            score, status = f"{e.score}/3", "✓ scored"
        lines.append(f"| {test} | {side} | {score} | {status} |")
    return "\n".join(lines)


def _finish_session(session_state):
    """Build the composite report + PDF for the whole session."""
    if not session_state or not session_state.entries:
        return ("⚠️ No clips analysed yet — analyse at least one clip first.",
                None, None, None)

    report, pdf_path = session_mod.finish_session(session_state)
    if report is None:
        return ("⚠️ Nothing to report.", None, None, None)

    if report.composite is not None:
        summary = [f"## Composite: {report.composite} / 21"]
    else:
        n = len(session_state.entries)
        summary = [f"## Composite: Incomplete — {n}/7 tests scored",
                   "*(One or more tests need clinician review or were unscored.)*"]

    if report.asymmetries:
        summary.append("\n### Asymmetries")
        for a in report.asymmetries:
            test = a["test"].replace("_", " ").title()
            summary.append(f"- **{test}:** L={a['left_score']} R={a['right_score']} (Δ {a['delta']})")

    flags = list(report.low_confidence_flags) + list(report.disagreement_flags)
    if flags:
        summary.append("\n### Flags")
        for fl in flags:
            summary.append(f"- {fl}")

    md_path = os.path.join(session_state.session_dir, "analysis.md")
    md_out = md_path if os.path.exists(md_path) else None
    return "\n".join(summary), pdf_path, md_out, report.scoresheet_path


# ─── App Builder ─────────────────────────────────────────────────────────────

def build_app() -> gr.Blocks:
    """Build the FormScout Gradio app with custom scout/trail theme."""
    with gr.Blocks(title="FormScout — FMS Screening Aid") as app:

        session_state = gr.State(None)

        # Header
        gr.HTML("""
        <div class="formscout-header">
            <h1>🏔️ FormScout</h1>
            <p style="color: #4a5f57; font-size: 0.95em;">
                Functional Movement Screen · Automated Scoring Aid
            </p>
        </div>
        """)

        # Safety banner (always visible — non-negotiable)
        gr.HTML(f'<div class="safety-banner">{DISCLAIMER}</div>')

        with gr.Row(equal_height=False):
            # Left column: Input
            with gr.Column(scale=2):
                gr.Markdown("### 📹 Input")
                video_input = gr.Video(label="Upload FMS Video")

                with gr.Row():
                    test_dropdown = gr.Dropdown(
                        choices=[name for name, _ in FMS_TESTS],
                        value="Deep Squat",
                        label="FMS Test",
                        scale=2,
                    )
                    side_dropdown = gr.Dropdown(
                        choices=["N/A", "Left", "Right"],
                        value="N/A",
                        label="Side",
                        scale=1,
                    )

                movement_guide = gr.Image(
                    value=_screening_image("Deep Squat"),
                    label="Movement guide",
                    interactive=False,
                    buttons=[],  # no download/share/fullscreen overlay
                    height=240,
                    elem_id="movement-guide",
                )

                _available_models = config.available_pose_models() or config.POSE_MODELS
                _default_model = (
                    config.DEFAULT_POSE_MODEL
                    if config.DEFAULT_POSE_MODEL in _available_models
                    else list(_available_models.keys())[0]
                )
                pose_model_dropdown = gr.Dropdown(
                    choices=list(_available_models.keys()),
                    value=_default_model,
                    label="Pose Model",
                )

                overlay_layers = gr.CheckboxGroup(
                    choices=["Skeleton", "Trails", "Velocity arrows"],
                    value=["Skeleton", "Trails", "Velocity arrows"],
                    label="Overlay Layers",
                )

                submit_btn = gr.Button(
                    "🎯 Score Movement",
                    variant="primary",
                    size="lg",
                )
                with gr.Row():
                    new_clip_btn = gr.Button("➕ Analyse new clip", visible=False)
                    finish_btn = gr.Button("✅ Finish & generate PDF",
                                           variant="primary", visible=False)

                gr.Markdown(
                    "*Tip: Film from the side at hip height for best accuracy. "
                    "One athlete, one rep per clip.*",
                    elem_classes=["topo-accent"],
                )

            # Right column: Results
            with gr.Column(scale=3):
                gr.Markdown("### 📊 Results")

                # Score display
                score_html = gr.HTML(value=_render_empty_state())

                # Tabs for details
                with gr.Tabs():
                    with gr.TabItem("📐 Rubric Breakdown"):
                        score_details = gr.Markdown("")

                    with gr.TabItem("🔧 Pipeline"):
                        pipeline_md = gr.Markdown("*Waiting for video...*")

                    with gr.TabItem("⚠️ Alerts"):
                        alerts_md = gr.Markdown("")

                    with gr.TabItem("🎬 Overlay Video"):
                        overlay_video = gr.Video(label="Annotated Movement")
                        velocity_md = gr.Markdown("")

                    with gr.TabItem("🗂️ Session"):
                        session_table = gr.Markdown("*No clips analysed yet.*")
                        finish_summary = gr.Markdown("")
                        scoresheet_file = gr.File(label="FMS Scoring Sheet (PDF)", visible=True)
                        pdf_file = gr.File(label="Detailed Screening Report (PDF)", visible=True)
                        md_file = gr.File(label="Analysis Log (Markdown)", visible=True)

        # Footer safety banner
        gr.HTML(f'<div class="safety-banner" style="margin-top: 20px;">{DISCLAIMER}</div>')

        gr.Markdown(
            "<center style='color: #6b7d75; font-size: 0.8em; margin-top: 12px;'>"
            "FormScout · ~18B params · Off the Grid · "
            "<a href='https://silastherapy.sk' style='color: #1f6e6e;'>Silas Therapy · Build Small Hackathon</a>"
            "</center>"
        )

        # ─── Event wiring ────────────────────────────────────────────────────

        def _map_inputs(video, test_display_name, side_display, pose_model_key, overlay_layers, sess):
            """Map UI display values to internal values and accumulate into the session."""
            test_map = {name: val for name, val in FMS_TESTS}
            test_name = test_map.get(test_display_name, "deep_squat")
            side = {"N/A": "na", "Left": "left", "Right": "right"}.get(side_display, "na")
            return process_video(video, test_name, side, pose_model_key, overlay_layers, sess)

        # Swap the movement-guide illustration when the selected test changes.
        test_dropdown.change(
            fn=_screening_image, inputs=[test_dropdown], outputs=[movement_guide],
        )

        submit_btn.click(
            fn=_map_inputs,
            inputs=[video_input, test_dropdown, side_dropdown, pose_model_dropdown,
                    overlay_layers, session_state],
            outputs=[session_state, score_html, pipeline_md, score_details, alerts_md,
                     overlay_video, velocity_md, session_table, new_clip_btn, finish_btn],
        )

        def _new_clip():
            """Clear inputs for the next clip; keep the session intact."""
            return None, _render_empty_state(), ""

        new_clip_btn.click(
            fn=_new_clip,
            inputs=[],
            outputs=[video_input, score_html, score_details],
        )

        finish_btn.click(
            fn=_finish_session,
            inputs=[session_state],
            outputs=[finish_summary, pdf_file, md_file, scoresheet_file],
        )

    return app


if __name__ == "__main__":
    app = build_app()
    app.launch(theme=formscout_theme(), css=FORMSCOUT_CSS,
               allowed_paths=[_ASSET_DIR])