training_coach/engine.py

from datetime import date
import logging
from math import ceil
import re
from collections.abc import Sequence

from training_coach.models import (
    CheckIn,
    CompletedSession,
    CompletedSet,
    Muscle,
    PlannedExercise,
    PrescribedSet,
    SessionPlan,
)


MVP_LOAD_INCREMENT_KG = 1.0
MVP_BASE_TARGET_RIR = 2
MVP_REP_STEP = 2
logger = logging.getLogger(__name__)

INJURY_KEYWORDS = (
    "ache",
    "aches",
    "hurt",
    "hurts",
    "injured",
    "injury",
    "pain",
    "painful",
    "pulled",
    "strain",
    "strained",
    "tear",
    "tearing",
    "tore",
    "torn",
)

MUSCLE_TEXT_MAP = {
    "bicep": Muscle.BICEPS_BRACHII,
    "biceps": Muscle.BICEPS_BRACHII,
    "tricep": Muscle.TRICEPS_BRACHII,
    "triceps": Muscle.TRICEPS_BRACHII,
    "hamstring": Muscle.HAMSTRINGS,
    "hamstrings": Muscle.HAMSTRINGS,
    "calf": Muscle.GASTROCNEMIUS,
    "calves": Muscle.GASTROCNEMIUS,
    "front delt": Muscle.FRONT_DELTOID,
    "front shoulder": Muscle.FRONT_DELTOID,
}


EXERCISE_MUSCLE_MAP = {
    "dumbbell-pullover": {
        Muscle.LATISSIMUS_DORSI,
        Muscle.PECTORALIS_MAJOR,
        Muscle.TRICEPS_BRACHII,
    },
    "dumbbell-row": {
        Muscle.LATISSIMUS_DORSI,
        Muscle.RHOMBOIDS,
        Muscle.TRAPEZIUS,
        Muscle.REAR_DELTOID,
        Muscle.BICEPS_BRACHII,
    },
    "barbell-incline-bench-press": {
        Muscle.PECTORALIS_MAJOR,
        Muscle.FRONT_DELTOID,
        Muscle.TRICEPS_BRACHII,
    },
    "dumbbell-incline-chest-fly": {
        Muscle.PECTORALIS_MAJOR,
        Muscle.FRONT_DELTOID,
    },
    "goblet-squat": {
        Muscle.QUADRICEPS,
        Muscle.GLUTEUS_MAXIMUS,
        Muscle.ADDUCTORS,
    },
    "barbell-skullcrusher": {Muscle.TRICEPS_BRACHII},
    "dumbbell-lateral-raise": {Muscle.SIDE_DELTOID},
    "barbell-triceps-extension": {Muscle.TRICEPS_BRACHII},
    "cable-lateral-raise": {Muscle.SIDE_DELTOID},
    "barbell-bicep-curl": {
        Muscle.BICEPS_BRACHII,
        Muscle.BRACHIALIS,
        Muscle.FOREARM_FLEXORS,
    },
    "barbell-hip-thrust": {
        Muscle.GLUTEUS_MAXIMUS,
        Muscle.HAMSTRINGS,
    },
    "barbell-standing-calf-raise": {
        Muscle.GASTROCNEMIUS,
        Muscle.SOLEUS,
    },
    "barbell-romanian-deadlift": {
        Muscle.HAMSTRINGS,
        Muscle.GLUTEUS_MAXIMUS,
        Muscle.SPINAL_ERECTORS,
    },
    "ez-bar-biceps-curl": {
        Muscle.BICEPS_BRACHII,
        Muscle.BRACHIALIS,
        Muscle.FOREARM_FLEXORS,
    },
    "cross-body-hammer-curl": {
        Muscle.BRACHIALIS,
        Muscle.BICEPS_BRACHII,
        Muscle.FOREARM_FLEXORS,
    },
}


def _sets(
    count: int,
    reps_low: int,
    reps_high: int | None = None,
    target_reps: int | None = None,
) -> list[PrescribedSet]:
    high = reps_high if reps_high is not None else reps_low
    return [
        PrescribedSet(
            set_number=set_number,
            target_reps_low=reps_low,
            target_reps_high=high,
            target_reps=target_reps,
        )
        for set_number in range(1, count + 1)
    ]


def _template_reps(item: dict) -> tuple[int, int, int | None]:
    if "reps" in item:
        target_reps = item["reps"]
        return max(1, target_reps - MVP_REP_STEP), target_reps + MVP_REP_STEP, target_reps

    reps_low = item["reps_low"]
    reps_high = item["reps_high"]
    return reps_low, reps_high, None


MVP_4_DAY_TEMPLATE = {
    1: [
        {
            "exercise_id": "dumbbell-pullover",
            "sets": 3,
            "reps": 13,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "dumbbell-row",
            "sets": 3,
            "reps": 16,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-incline-bench-press",
            "sets": 3,
            "reps": 10,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "dumbbell-incline-chest-fly",
            "sets": 3,
            "reps": 12,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "goblet-squat",
            "sets": 4,
            "reps": 12,
            "rest_seconds": 60,
        },
    ],
    2: [
        {
            "exercise_id": "barbell-skullcrusher",
            "sets": 3,
            "reps": 8,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "dumbbell-lateral-raise",
            "sets": 3,
            "reps": 20,
            "rest_seconds": 60,
            "notes": "Lean-away DB lateral raise.",
        },
        {
            "exercise_id": "barbell-triceps-extension",
            "sets": 3,
            "reps_low": 10,
            "reps_high": 12,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "cable-lateral-raise",
            "sets": 3,
            "reps": 15,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-bicep-curl",
            "sets": 5,
            "reps": 10,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-hip-thrust",
            "sets": 5,
            "reps": 10,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-standing-calf-raise",
            "sets": 4,
            "reps": 10,
            "rest_seconds": 60,
        },
    ],
    3: [
        {
            "exercise_id": "barbell-incline-bench-press",
            "sets": 3,
            "reps": 8,
            "rest_seconds": 60,
            "notes": "Wide grip.",
        },
        {
            "exercise_id": "dumbbell-incline-chest-fly",
            "sets": 3,
            "reps": 10,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "dumbbell-row",
            "sets": 3,
            "reps": 24,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "dumbbell-pullover",
            "sets": 3,
            "reps": 15,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-romanian-deadlift",
            "sets": 4,
            "reps": 10,
            "rest_seconds": 60,
        },
    ],
    4: [
        {
            "exercise_id": "ez-bar-biceps-curl",
            "sets": 3,
            "reps": 12,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "dumbbell-lateral-raise",
            "sets": 3,
            "reps": 20,
            "rest_seconds": 60,
            "notes": "Lean away.",
        },
        {
            "exercise_id": "cross-body-hammer-curl",
            "sets": 3,
            "reps": 18,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "cable-lateral-raise",
            "sets": 3,
            "reps": 18,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-triceps-extension",
            "sets": 3,
            "reps": 12,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "goblet-squat",
            "sets": 4,
            "reps": 12,
            "rest_seconds": 60,
        },
        {
            "exercise_id": "barbell-standing-calf-raise",
            "sets": 4,
            "reps": 10,
            "rest_seconds": 60,
        },
    ],
}


def build_session_for_day(
    day_number: int,
    session_date: date,
    check_in: CheckIn,
    completed_sessions: Sequence[CompletedSession] | None = None,
) -> SessionPlan:
    if day_number not in MVP_4_DAY_TEMPLATE:
        raise ValueError("day_number must be between 1 and 4")

    logger.info(
        "event=engine_build_start day_number=%s history_sessions=%s "
        "time_minutes=%s pain_or_injury=%s pain_issues=%s",
        day_number,
        len(completed_sessions or []),
        check_in.time_available_minutes,
        check_in.pain_or_injury,
        len(check_in.pain_issues),
    )
    planned_exercises = []
    for order, item in enumerate(MVP_4_DAY_TEMPLATE[day_number], start=1):
        reps_low, reps_high, target_reps = _template_reps(item)
        planned_exercises.append(
            PlannedExercise(
                exercise_id=item["exercise_id"],
                order=order,
                prescribed_sets=_sets(item["sets"], reps_low, reps_high, target_reps),
                rest_seconds=item["rest_seconds"],
                notes=item.get("notes", ""),
            )
        )

    template_set_count = sum(
        len(exercise.prescribed_sets)
        for exercise in planned_exercises
    )
    logger.info(
        "event=engine_template_loaded day_number=%s exercises=%s sets=%s",
        day_number,
        len(planned_exercises),
        template_set_count,
    )

    if completed_sessions:
        planned_exercises = [
            apply_double_progression(exercise, completed_sessions)
            for exercise in planned_exercises
        ]
        logger.info(
            "event=engine_progression_applied day_number=%s exercises=%s",
            day_number,
            len(planned_exercises),
        )

    before_pain_count = len(planned_exercises)
    planned_exercises, pain_filter_notes = apply_pain_filter(
        planned_exercises,
        check_in,
    )
    if pain_filter_notes:
        logger.info(
            "event=engine_pain_filter_applied before_exercises=%s after_exercises=%s",
            before_pain_count,
            len(planned_exercises),
        )

    session_notes = [
        f"MVP fixed day {day_number} template.",
        "Double progression applies +1 kg after all sets hit the top rep target.",
    ]
    if pain_filter_notes:
        session_notes.extend(pain_filter_notes)

    before_time_sets = sum(
        len(exercise.prescribed_sets)
        for exercise in planned_exercises
    )
    planned_exercises, time_compression_notes = apply_time_compression(
        planned_exercises,
        check_in,
    )
    if time_compression_notes:
        session_notes.extend(time_compression_notes)
        after_time_sets = sum(
            len(exercise.prescribed_sets)
            for exercise in planned_exercises
        )
        logger.info(
            "event=engine_time_compression_applied before_sets=%s after_sets=%s time_minutes=%s",
            before_time_sets,
            after_time_sets,
            check_in.time_available_minutes,
        )

    before_readiness_sets = sum(
        len(exercise.prescribed_sets)
        for exercise in planned_exercises
    )
    planned_exercises, readiness_notes = apply_readiness_modifier(
        planned_exercises,
        check_in,
    )
    if readiness_notes:
        session_notes.extend(readiness_notes)
        after_readiness_sets = sum(
            len(exercise.prescribed_sets)
            for exercise in planned_exercises
        )
        logger.info(
            "event=engine_readiness_applied before_sets=%s after_sets=%s",
            before_readiness_sets,
            after_readiness_sets,
        )

    final_set_count = sum(
        len(exercise.prescribed_sets)
        for exercise in planned_exercises
    )
    logger.info(
        "event=engine_build_complete day_number=%s exercises=%s sets=%s notes=%s",
        day_number,
        len(planned_exercises),
        final_set_count,
        len(session_notes),
    )

    return SessionPlan(
        date=session_date,
        check_in=check_in,
        planned_exercises=planned_exercises,
        notes=" ".join(session_notes),
    )


def next_day_after(day_number: int) -> int:
    if day_number not in MVP_4_DAY_TEMPLATE:
        raise ValueError("day_number must be between 1 and 4")
    return 1 if day_number == 4 else day_number + 1


def suggest_next_training_day(completed_sessions: Sequence[CompletedSession]) -> int:
    if not completed_sessions:
        return 1
    return next_day_after(completed_sessions[-1].day_number)


def _latest_completed_sets_for_exercise(
    completed_sessions: Sequence[CompletedSession],
    exercise_id: str,
) -> list[CompletedSet]:
    for session in reversed(completed_sessions):
        sets = [
            completed_set
            for completed_set in session.completed_sets
            if completed_set.exercise_id == exercise_id
        ]
        if sets:
            return sorted(sets, key=lambda completed_set: completed_set.set_number)
    return []


def _hit_top_reps(
    planned_exercise: PlannedExercise,
    completed_sets: Sequence[CompletedSet],
) -> bool:
    completed_by_number = {
        completed_set.set_number: completed_set for completed_set in completed_sets
    }
    for prescribed_set in planned_exercise.prescribed_sets:
        completed_set = completed_by_number.get(prescribed_set.set_number)
        if completed_set is None:
            return False
        if completed_set.actual_reps < prescribed_set.target_reps_high:
            return False
    return True


def _next_target_reps(
    prescribed_set: PrescribedSet,
    completed_set: CompletedSet | None,
    should_increase_load: bool,
) -> int | None:
    if completed_set is None:
        return prescribed_set.target_reps
    if should_increase_load:
        return prescribed_set.target_reps_low
    return min(
        prescribed_set.target_reps_high,
        max(prescribed_set.target_reps_low, completed_set.actual_reps + MVP_REP_STEP),
    )


def apply_double_progression(
    planned_exercise: PlannedExercise,
    completed_sessions: Sequence[CompletedSession],
    load_increment_kg: float = MVP_LOAD_INCREMENT_KG,
) -> PlannedExercise:
    completed_sets = _latest_completed_sets_for_exercise(
        completed_sessions,
        planned_exercise.exercise_id,
    )
    if not completed_sets:
        return planned_exercise

    completed_by_number = {
        completed_set.set_number: completed_set for completed_set in completed_sets
    }
    should_increase_load = _hit_top_reps(planned_exercise, completed_sets)
    prescribed_sets = []
    for prescribed_set in planned_exercise.prescribed_sets:
        completed_set = completed_by_number.get(prescribed_set.set_number)
        target_load = None
        target_reps = _next_target_reps(
            prescribed_set,
            completed_set,
            should_increase_load,
        )
        if completed_set is not None:
            target_load = completed_set.actual_load
            if should_increase_load:
                target_load += load_increment_kg

        prescribed_sets.append(
            prescribed_set.model_copy(
                update={
                    "target_load": target_load,
                    "target_reps": target_reps,
                }
            )
        )

    progression_note = (
        "Progression: top of range hit; add 1 kg and reset reps to the low end."
        if should_increase_load
        else f"Progression: repeat latest load and add up to {MVP_REP_STEP} reps."
    )
    notes = (
        f"{planned_exercise.notes} {progression_note}".strip()
        if planned_exercise.notes
        else progression_note
    )
    return planned_exercise.model_copy(
        update={"prescribed_sets": prescribed_sets, "notes": notes}
    )


def _pain_muscles(check_in: CheckIn) -> set[Muscle]:
    explicit_muscles = {
        issue.affected_muscle
        for issue in check_in.pain_issues
        if issue.affected_muscle is not None
    }
    manual_text = f"{check_in.raw_text} {check_in.soreness}"
    inferred_muscles = set()
    if check_in.pain_or_injury == "yes" or _mentions_injury(manual_text):
        inferred_muscles = _infer_muscles_from_text(manual_text)

    if check_in.pain_or_injury != "yes" and not _mentions_injury(manual_text):
        return set()

    return explicit_muscles.union(inferred_muscles)


def _contains_word(text: str, word: str) -> bool:
    return re.search(rf"\b{re.escape(word)}\b", text.lower()) is not None


def _mentions_injury(text: str) -> bool:
    return any(_contains_word(text, keyword) for keyword in INJURY_KEYWORDS)


def _infer_muscles_from_text(text: str) -> set[Muscle]:
    normalized = text.lower()
    return {
        muscle
        for phrase, muscle in MUSCLE_TEXT_MAP.items()
        if phrase in normalized
    }


def apply_pain_filter(
    planned_exercises: Sequence[PlannedExercise],
    check_in: CheckIn,
) -> tuple[list[PlannedExercise], list[str]]:
    painful_muscles = _pain_muscles(check_in)
    if not painful_muscles:
        return list(planned_exercises), []

    kept_exercises = []
    removed_exercise_ids = []
    for planned_exercise in planned_exercises:
        exercise_muscles = EXERCISE_MUSCLE_MAP.get(planned_exercise.exercise_id, set())
        if exercise_muscles.intersection(painful_muscles):
            removed_exercise_ids.append(planned_exercise.exercise_id)
            continue
        kept_exercises.append(planned_exercise)

    reordered_exercises = [
        exercise.model_copy(update={"order": order})
        for order, exercise in enumerate(kept_exercises, start=1)
    ]
    if not removed_exercise_ids:
        return reordered_exercises, []

    removed_text = ", ".join(removed_exercise_ids)
    muscles_text = ", ".join(sorted(muscle.value for muscle in painful_muscles))
    return reordered_exercises, [
        f"Pain filter removed {removed_text} because of affected muscle(s): {muscles_text}."
    ]


def _copy_first_sets(planned_exercise: PlannedExercise, set_count: int) -> PlannedExercise:
    prescribed_sets = [
        prescribed_set.model_copy(update={"set_number": set_number})
        for set_number, prescribed_set in enumerate(
            planned_exercise.prescribed_sets[:set_count],
            start=1,
        )
    ]
    return planned_exercise.model_copy(update={"prescribed_sets": prescribed_sets})


def _compress_to_target_sets(
    planned_exercises: Sequence[PlannedExercise],
    target_set_count: int,
    minimum_sets_per_kept_exercise: int,
) -> list[PlannedExercise]:
    set_counts = [len(exercise.prescribed_sets) for exercise in planned_exercises]
    current_set_count = sum(set_counts)

    for index in range(len(set_counts) - 1, -1, -1):
        while (
            current_set_count > target_set_count
            and set_counts[index] > minimum_sets_per_kept_exercise
        ):
            set_counts[index] -= 1
            current_set_count -= 1

    for index in range(len(set_counts) - 1, -1, -1):
        while current_set_count > target_set_count and set_counts[index] > 0:
            set_counts[index] -= 1
            current_set_count -= 1

    compressed = []
    for exercise, set_count in zip(planned_exercises, set_counts):
        if set_count <= 0:
            continue
        compressed.append(_copy_first_sets(exercise, set_count))

    return [
        exercise.model_copy(update={"order": order})
        for order, exercise in enumerate(compressed, start=1)
    ]


def apply_time_compression(
    planned_exercises: Sequence[PlannedExercise],
    check_in: CheckIn,
) -> tuple[list[PlannedExercise], list[str]]:
    available_minutes = check_in.time_available_minutes
    if available_minutes is None or available_minutes >= 60:
        return list(planned_exercises), []

    total_sets = sum(len(exercise.prescribed_sets) for exercise in planned_exercises)
    if available_minutes < 30:
        candidate_exercises = list(planned_exercises[:4])
        target_set_count = min(8, sum(len(exercise.prescribed_sets) for exercise in candidate_exercises))
        compressed = _compress_to_target_sets(
            candidate_exercises,
            target_set_count=target_set_count,
            minimum_sets_per_kept_exercise=1,
        )
        return compressed, [
            "Time compression: under 30 minutes, kept the first four exercises and capped work at 8 sets."
        ]

    if available_minutes < 45:
        target_set_count = max(1, ceil(total_sets * 0.60))
        compressed = _compress_to_target_sets(
            planned_exercises,
            target_set_count=target_set_count,
            minimum_sets_per_kept_exercise=1,
        )
        return compressed, [
            f"Time compression: {available_minutes} minutes, reduced planned sets from {total_sets} to {target_set_count}."
        ]

    target_set_count = max(1, ceil(total_sets * 0.75))
    compressed = _compress_to_target_sets(
        planned_exercises,
        target_set_count=target_set_count,
        minimum_sets_per_kept_exercise=2,
    )
    return compressed, [
        f"Time compression: {available_minutes} minutes, reduced planned sets from {total_sets} to {target_set_count}."
    ]


def _sleep_quality_score(check_in: CheckIn) -> int:
    return {
        "poor": 1,
        "okay": 3,
        "good": 5,
        None: 3,
    }[check_in.sleep_quality]


def _sleep_duration_score(check_in: CheckIn) -> int:
    hours = check_in.sleep_hours
    if hours is None:
        return 3
    if hours < 5:
        return 1
    if hours < 6:
        return 2
    if hours < 7:
        return 3
    if hours <= 8.5:
        return 4
    return 5


def _energy_score(check_in: CheckIn) -> int:
    return {
        "low": 1,
        "medium": 3,
        "high": 5,
        None: 3,
    }[check_in.energy_level]


def _soreness_score(check_in: CheckIn) -> int:
    soreness = check_in.soreness.lower()
    if any(word in soreness for word in ("severe", "extreme", "very sore")):
        return 1
    if any(word in soreness for word in ("sore", "tight", "stiff", "ache")):
        return 2
    if any(phrase in soreness for phrase in ("no soreness", "not sore", "none")):
        return 5
    return 3


def _mood_stress_score(check_in: CheckIn) -> int:
    return {
        "stressed": 1,
        "neutral": 3,
        "ready": 5,
        None: 3,
    }[check_in.mood_stress]


def readiness_score(check_in: CheckIn) -> float:
    return (
        _sleep_quality_score(check_in) * 0.20
        + _sleep_duration_score(check_in) * 0.15
        + _energy_score(check_in) * 0.25
        + _soreness_score(check_in) * 0.15
        + _mood_stress_score(check_in) * 0.15
        + _mood_stress_score(check_in) * 0.10
    )


def _readiness_modifier(check_in: CheckIn) -> tuple[float, int, str]:
    score = readiness_score(check_in)
    if score < 2.5:
        return 0.50, 2, f"Readiness: very low ({score:.1f}/5), reduced sets by about 50% and added +2 RIR."
    if score < 3.0:
        return 0.80, 1, f"Readiness: low ({score:.1f}/5), reduced sets by about 20% and added +1 RIR."
    if score > 4.2:
        return 1.00, 0, f"Readiness: high ({score:.1f}/5), MVP keeps the plan unchanged."
    return 1.00, 0, f"Readiness: normal ({score:.1f}/5), plan unchanged."


def _with_target_rir(planned_exercises: Sequence[PlannedExercise], target_rir: int) -> list[PlannedExercise]:
    exercises = []
    for exercise in planned_exercises:
        prescribed_sets = [
            prescribed_set.model_copy(update={"target_rir": target_rir})
            for prescribed_set in exercise.prescribed_sets
        ]
        exercises.append(exercise.model_copy(update={"prescribed_sets": prescribed_sets}))
    return exercises


def apply_readiness_modifier(
    planned_exercises: Sequence[PlannedExercise],
    check_in: CheckIn,
) -> tuple[list[PlannedExercise], list[str]]:
    set_multiplier, rir_delta, note = _readiness_modifier(check_in)
    target_rir = MVP_BASE_TARGET_RIR + rir_delta

    if check_in.sleep_hours is not None and check_in.sleep_hours < 5:
        set_multiplier = min(set_multiplier, 0.75)
        target_rir = max(target_rir, 3)
        note = f"{note} Sleep override: under 5 hours, capped stress at minimum RIR 3."

    total_sets = sum(len(exercise.prescribed_sets) for exercise in planned_exercises)
    if set_multiplier < 1:
        target_set_count = max(1, ceil(total_sets * set_multiplier))
        planned_exercises = _compress_to_target_sets(
            planned_exercises,
            target_set_count=target_set_count,
            minimum_sets_per_kept_exercise=1,
        )

    planned_exercises = _with_target_rir(planned_exercises, target_rir)
    return planned_exercises, [note]