training_coach/models.py

from enum import StrEnum
from datetime import date
from typing import Literal

from pydantic import BaseModel, ConfigDict, Field, model_validator


EnergyLevel = Literal["low", "medium", "high"]
SleepQuality = Literal["poor", "okay", "good"]
PainOrInjury = Literal["yes", "no", "unsure"]
PainSeverity = Literal["mild", "moderate", "severe", "unsure"]
MoodStress = Literal["stressed", "neutral", "ready"]
ContextSignalLabel = Literal[
    "recent_endurance_run",
    "travel",
    "illness",
    "alcohol",
    "unusual_stress",
    "unusually_hard_work",
    "other",
]
FollowUpField = Literal[
    "time_available_minutes",
    "energy_level",
    "sleep_quality",
    "sleep_hours",
    "soreness",
    "pain_or_injury",
    "pain_issues",
    "mood_stress",
    "context_signals",
    "readiness",
    "other",
]


class CoachModel(BaseModel):
    model_config = ConfigDict(extra="forbid")


class Muscle(StrEnum):
    PECTORALIS_MAJOR = "pectoralis_major"
    LATISSIMUS_DORSI = "latissimus_dorsi"
    TRAPEZIUS = "trapezius"
    RHOMBOIDS = "rhomboids"
    FRONT_DELTOID = "front_deltoid"
    SIDE_DELTOID = "side_deltoid"
    REAR_DELTOID = "rear_deltoid"
    BICEPS_BRACHII = "biceps_brachii"
    BRACHIALIS = "brachialis"
    TRICEPS_BRACHII = "triceps_brachii"
    FOREARM_FLEXORS = "forearm_flexors"
    FOREARM_EXTENSORS = "forearm_extensors"
    RECTUS_ABDOMINIS = "rectus_abdominis"
    OBLIQUES = "obliques"
    SPINAL_ERECTORS = "spinal_erectors"
    GLUTEUS_MAXIMUS = "gluteus_maximus"
    GLUTEUS_MEDIUS = "gluteus_medius"
    QUADRICEPS = "quadriceps"
    HAMSTRINGS = "hamstrings"
    ADDUCTORS = "adductors"
    ABDUCTORS = "abductors"
    GASTROCNEMIUS = "gastrocnemius"
    SOLEUS = "soleus"


class CheckIn(CoachModel):
    raw_text: str = ""
    time_available_minutes: int | None = Field(default=None, gt=0)
    energy_level: EnergyLevel | None = None
    sleep_quality: SleepQuality | None = None
    sleep_hours: float | None = Field(default=None, ge=0, le=24)
    soreness: str = ""
    pain_or_injury: PainOrInjury = "unsure"
    pain_issues: list["PainIssue"] = Field(default_factory=list)
    mood_stress: MoodStress | None = None
    notes: str = ""


class PainIssue(CoachModel):
    affected_muscle: Muscle | None = None
    severity: PainSeverity = "unsure"
    notes: str = Field(min_length=1)


class ContextSignal(CoachModel):
    label: ContextSignalLabel
    evidence: str = Field(min_length=1)
    follow_up_question: str = ""


class FollowUpQuestion(CoachModel):
    field: FollowUpField
    question: str = Field(min_length=1)
    reason: str = Field(min_length=1)


class ParsedCheckIn(CoachModel):
    check_in: CheckIn
    missing_fields: list[str] = Field(default_factory=list)
    follow_up_items: list[FollowUpQuestion] = Field(default_factory=list)
    follow_up_questions: list[str] = Field(default_factory=list)
    context_signals: list[ContextSignal] = Field(default_factory=list)
    notes: str = ""


class Exercise(CoachModel):
    id: str = Field(min_length=1)
    name: str = Field(min_length=1)
    primary_muscle_group: Muscle
    secondary_muscle_groups: list[Muscle] = Field(default_factory=list)
    is_compound: bool
    rep_range_low: int = Field(gt=0)
    rep_range_high: int = Field(gt=0)

    @model_validator(mode="after")
    def validate_rep_range(self):
        if self.rep_range_low > self.rep_range_high:
            raise ValueError("rep_range_low must be less than or equal to rep_range_high")
        return self


class PrescribedSet(CoachModel):
    set_number: int = Field(gt=0)
    target_reps_low: int = Field(gt=0)
    target_reps_high: int = Field(gt=0)
    target_reps: int | None = Field(default=None, gt=0)
    target_load: float | None = Field(default=None, ge=0)
    target_rir: int | None = Field(default=None, ge=0)

    @model_validator(mode="after")
    def validate_target_rep_range(self):
        if self.target_reps_low > self.target_reps_high:
            raise ValueError("target_reps_low must be less than or equal to target_reps_high")
        if self.target_reps is not None and not (
            self.target_reps_low <= self.target_reps <= self.target_reps_high
        ):
            raise ValueError("target_reps must be inside the target rep range")
        return self


class PlannedExercise(CoachModel):
    exercise_id: str = Field(min_length=1)
    order: int = Field(gt=0)
    prescribed_sets: list[PrescribedSet] = Field(min_length=1)
    rest_seconds: int | None = Field(default=None, ge=0)
    notes: str = ""


class SessionPlan(CoachModel):
    date: date
    check_in: CheckIn
    planned_exercises: list[PlannedExercise] = Field(min_length=1)
    notes: str = ""


class LoggedSet(CoachModel):
    set_number: int = Field(gt=0)
    actual_reps: int = Field(gt=0)
    actual_load: float = Field(ge=0)
    rpe: float | None = Field(default=None, ge=0, le=10)
    rest_seconds_before: int | None = Field(default=None, ge=0)
    notes: str = ""


class LoggedExercise(CoachModel):
    exercise_id: str = Field(min_length=1)
    order: int = Field(gt=0)
    logged_sets: list[LoggedSet] = Field(min_length=1)
    notes: str = ""


class SessionLog(CoachModel):
    date: date
    planned_session: SessionPlan
    logged_exercises: list[LoggedExercise] = Field(min_length=1)
    notes: str = ""


class CompletedSet(CoachModel):
    exercise_id: str = Field(min_length=1)
    set_number: int = Field(gt=0)
    actual_reps: int = Field(gt=0)
    actual_load: float = Field(ge=0)
    rpe: float | None = Field(default=None, ge=0, le=10)
    notes: str = ""


class CompletedSession(CoachModel):
    date: date
    day_number: int = Field(ge=1, le=4)
    completed_sets: list[CompletedSet] = Field(min_length=1)
    notes: str = ""