main.py

# main.py
# Post-Workout FastAPI app
# Designed to be called from a SwiftUI app

import os
import anthropic
from contextlib import asynccontextmanager
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
from models.model import (
    load_models,
    run_inference,
    PHYSICAL_LABEL_COLS, PHYSICAL_DECODERS,
    MENTAL_LABEL_COLS,   MENTAL_DECODERS,
)

# ─────────────────────────────────────────────
# STARTUP — load both models once into memory
# ─────────────────────────────────────────────

ml_models = {}

@asynccontextmanager
async def lifespan(app: FastAPI):
    tokenizer, physical_model, mental_model = load_models()
    ml_models["tokenizer"]      = tokenizer
    ml_models["physical_model"] = physical_model
    ml_models["mental_model"]   = mental_model
    print("All models loaded and ready")
    yield
    ml_models.clear()


# ─────────────────────────────────────────────
# APP SETUP
# ─────────────────────────────────────────────

app = FastAPI(
    title="Post Workout API",
    description="Classifies post-workout physical and mental state, returns a Claude coaching summary and tips.",
    version="1.0.0",
    lifespan=lifespan
)

# CORS — allows SwiftUI app to hit the API from any origin
# Tighten origins to your app's domain once in production
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["POST", "GET"],
    allow_headers=["*"],
)


# ─────────────────────────────────────────────
# REQUEST / RESPONSE SCHEMAS
# ─────────────────────────────────────────────

class UserProfile(BaseModel):
    primary_goal: str           # e.g. "muscle_hypertrophy"
    modifiers:    list[str]     # e.g. ["home_based", "time_constrained"]


class PredictRequest(BaseModel):
    text:               str               # raw user post-workout text — passed to BERT only
    duration_minutes:   int               # actual session duration
    workout_type:       str               # e.g. "strength", "cardio", "mobility"
    next_session_timing: str              # e.g. "training_tomorrow", "competition_soon"
    user_profile:       UserProfile
    planned_workout:    Optional[str] = None  # Claude output from pre-workout coach, if used


class LabelResult(BaseModel):
    label:      str             # decoded string label e.g. "mild"
    confidence: float           # 0.0 – 1.0


class PhysicalOutput(BaseModel):
    pain_label:          LabelResult
    completion_label:    LabelResult
    fatigue_label:       LabelResult
    recovery_need_label: LabelResult


class MentalOutput(BaseModel):
    performance_label:  LabelResult
    satisfaction_label: LabelResult
    pr_achieved_label:  LabelResult
    motivation_label:   LabelResult


class PredictResponse(BaseModel):
    physical:       PhysicalOutput
    mental:         MentalOutput
    coach_response: str             # Claude's workout summary + tips
    enriched_text:  str             # text sent to BERT (for debugging)


# ─────────────────────────────────────────────
# PROMPT BUILDER
# ─────────────────────────────────────────────

def build_claude_prompt(goal, duration_minutes, workout_type, next_session_timing, physical, mental, planned_workout):
    planned_section = (
        f"Planned session   : {planned_workout}"
        if planned_workout
        else "Planned session   : none — user worked out independently"
    )

    return f"""
You are a personal fitness coach reviewing a completed training session with your athlete.

User goal          : {goal}
Session duration   : {duration_minutes} minutes
Workout type       : {workout_type}
Next session       : {next_session_timing}
{planned_section}

Physical state:
- Pain          : {physical['pain_label']['label']}
- Completion    : {physical['completion_label']['label']}
- Fatigue       : {physical['fatigue_label']['label']}
- Recovery need : {physical['recovery_need_label']['label']}

Mental state:
- Performance   : {mental['performance_label']['label']}
- Satisfaction  : {mental['satisfaction_label']['label']}
- PR achieved   : {mental['pr_achieved_label']['label']}
- Motivation    : {mental['motivation_label']['label']}

REASONING RULES — apply before writing:
1. If no planned session was provided, skip the completion comparison entirely.
   Do not reference or imply there was a plan. Focus only on what the labels
   reveal about how the session went.
2. PR handling: if pr_achieved is true, open with a genuine, specific celebration tied to the user's goal — not just "great job." Explain what this PR signals about their progress.
3. Recovery routing — cross-reference recovery_need with next_session_timing:
   - recovery_need = rest_day → tomorrow's advice must be strictly non-training: sleep, nutrition, or mobility only. Do not suggest any workout regardless of next_session_timing.
   - recovery_need = light_day → suggest one low-intensity option only (walk, yoga, or easy swim).
   - recovery_need = normal or ready → give a forward-looking tip shaped by next_session_timing:
       - training_tomorrow    → prioritize sleep and nutrition; flag any soreness that needs monitoring overnight.
       - two_days_rest        → can suggest slightly more demanding recovery work.
       - competition_soon     → tomorrow is strictly preparation only — no training suggestions; focus on peaking strategy.
       - deload_week          → all tips should reflect reduced intensity expectations; discourage pushing hard.
       - no_plan_yet          → give a general recovery recommendation without assuming timing.
4. Pain flag: if pain label is moderate or high, add a one-line note recommending the user monitor that area and consider professional advice if it persists.

OUTPUT FORMATTING RULES:
- Use # to prefix every section header e.g. "# Session verdict"
- Write each content line as plain text with no bullet or prefix — the app will handle bullets
- Use | to separate a line's main point from its supporting detail e.g. "Great squat depth | 3 sets completed clean"
- Leave a blank line between each section
- No markdown. No asterisks. No bold. No numbered lists.

OUTPUT FORMAT — respond in exactly this structure:
Session verdict: [one sentence — honest assessment of how the session went vs. what was planned]
Highlights: [one specific thing that went well, tied to their goal]
Tips for next time: [2-3 concrete, specific tips derived directly from today's labels — not generic advice]
Tomorrow: [one clear recommendation based on recovery_need and next_session_timing routing above]

Be direct and specific. Reference the planned session when relevant. Never give advice that could apply to any athlete — make it personal to today.
""".strip()


# ─────────────────────────────────────────────
# ENDPOINTS
# ─────────────────────────────────────────────

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    """
    Main endpoint — called by SwiftUI after user submits post-workout summary.

    Flow:
      1. Enriches user text with goal (for BERT only)
      2. Runs physical and mental BERT models concurrently
      3. Builds Claude prompt from labels + session metadata + optional planned workout
      4. Calls Claude API
      5. Returns classifications + coaching response to SwiftUI
    """
    try:
        # 1. Enrich text with goal for BERT — duration, workout_type, and
        #    planned_workout are kept separate so they don't pollute the BERT
        #    input distribution
        enriched_text = f"Goal: {request.user_profile.primary_goal}. {request.text}"

        # 2. Run physical and mental models concurrently
        physical_output, mental_output = await asyncio.gather(
            asyncio.to_thread(run_inference, ml_models["physical_model"], ml_models["tokenizer"], enriched_text, PHYSICAL_LABEL_COLS, PHYSICAL_DECODERS),
            asyncio.to_thread(run_inference, ml_models["mental_model"],   ml_models["tokenizer"], enriched_text, MENTAL_LABEL_COLS,   MENTAL_DECODERS)
        )

        # 3. Build Claude prompt with all available context
        prompt = build_claude_prompt(
            request.user_profile.primary_goal,
            request.duration_minutes,
            request.workout_type,
            request.next_session_timing,
            physical_output,
            mental_output,
            request.planned_workout
        )

        # 4. Call Claude API
        claude_client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])
        message       = claude_client.messages.create(
            model="claude-sonnet-4-20250514",
            max_tokens=1000,
            messages=[{"role": "user", "content": prompt}]
        )

        # 5. Return full response
        return PredictResponse(
            physical=physical_output,
            mental=mental_output,
            coach_response=message.content[0].text,
            enriched_text=enriched_text
        )

    except KeyError as e:
        raise HTTPException(status_code=500, detail=f"Model not loaded: {str(e)}")
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/health")
def health():
    """
    Health check endpoint.
    SwiftUI app can ping this after a workout to pre-warm the Space.
    """
    return {
        "status":         "ok",
        "models_loaded":  len(ml_models) == 3,
        "physical_ready": "physical_model" in ml_models,
        "mental_ready":   "mental_model"   in ml_models,
    }


@app.get("/")
def root():
    return {"message": "Post Workout API is running. POST to /predict to get started."}