tests/test_engine.py

from datetime import date
import unittest

from training_coach.engine import (
    EXERCISE_MUSCLE_MAP,
    MVP_LOAD_INCREMENT_KG,
    apply_pain_filter,
    apply_time_compression,
    apply_readiness_modifier,
    apply_double_progression,
    build_session_for_day,
    next_day_after,
    readiness_score,
    suggest_next_training_day,
)
from training_coach.models import CheckIn, CompletedSession, CompletedSet, Muscle, PainIssue


class EngineTemplateTest(unittest.TestCase):
    def test_day_1_preserves_exercise_order_and_prescription(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text="60 min, medium energy"),
        )

        self.assertEqual(len(plan.planned_exercises), 5)
        self.assertEqual(plan.planned_exercises[0].exercise_id, "dumbbell-pullover")
        self.assertEqual(plan.planned_exercises[0].order, 1)
        self.assertEqual(plan.planned_exercises[0].rest_seconds, 60)
        self.assertEqual(len(plan.planned_exercises[0].prescribed_sets), 3)
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_reps_low, 11)
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_reps_high, 15)
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_reps, 13)
        self.assertEqual(plan.planned_exercises[-1].exercise_id, "goblet-squat")
        self.assertEqual(len(plan.planned_exercises[-1].prescribed_sets), 4)

    def test_day_2_preserves_rep_range_and_notes_with_mvp_rest(self):
        plan = build_session_for_day(2, date(2026, 6, 11), CheckIn(raw_text=""))

        triceps_extension = plan.planned_exercises[2]
        self.assertEqual(triceps_extension.exercise_id, "barbell-triceps-extension")
        self.assertEqual(triceps_extension.rest_seconds, 60)
        self.assertEqual(triceps_extension.prescribed_sets[0].target_reps_low, 10)
        self.assertEqual(triceps_extension.prescribed_sets[0].target_reps_high, 12)
        self.assertIn("Lean-away", plan.planned_exercises[1].notes)

    def test_day_3_preserves_row_reps_with_mvp_rest(self):
        plan = build_session_for_day(3, date(2026, 6, 11), CheckIn(raw_text=""))

        dumbbell_row = plan.planned_exercises[2]
        self.assertEqual(dumbbell_row.exercise_id, "dumbbell-row")
        self.assertEqual(dumbbell_row.rest_seconds, 60)
        self.assertEqual(dumbbell_row.prescribed_sets[0].target_reps_low, 22)
        self.assertEqual(dumbbell_row.prescribed_sets[0].target_reps_high, 26)
        self.assertEqual(dumbbell_row.prescribed_sets[0].target_reps, 24)

    def test_day_4_uses_mvp_rest_for_goblet_squat(self):
        plan = build_session_for_day(4, date(2026, 6, 11), CheckIn(raw_text=""))

        goblet_squat = plan.planned_exercises[5]
        self.assertEqual(goblet_squat.exercise_id, "goblet-squat")
        self.assertEqual(goblet_squat.rest_seconds, 60)

    def test_all_template_exercises_use_one_minute_rest_for_mvp(self):
        for day_number in range(1, 5):
            plan = build_session_for_day(
                day_number,
                date(2026, 6, 11),
                CheckIn(raw_text=""),
            )
            self.assertTrue(
                all(exercise.rest_seconds == 60 for exercise in plan.planned_exercises)
            )

    def test_rejects_unknown_day(self):
        with self.assertRaises(ValueError):
            build_session_for_day(5, date(2026, 6, 11), CheckIn(raw_text=""))

    def test_next_day_rotates_through_four_day_template(self):
        self.assertEqual(next_day_after(1), 2)
        self.assertEqual(next_day_after(2), 3)
        self.assertEqual(next_day_after(3), 4)
        self.assertEqual(next_day_after(4), 1)

    def test_next_day_rejects_unknown_day(self):
        with self.assertRaises(ValueError):
            next_day_after(0)

    def test_suggest_next_training_day_starts_at_day_1_without_history(self):
        self.assertEqual(suggest_next_training_day([]), 1)

    def test_suggest_next_training_day_uses_last_completed_session(self):
        completed_set = CompletedSet(
            exercise_id="dumbbell-row",
            set_number=1,
            actual_reps=16,
            actual_load=32.5,
        )
        history = [
            CompletedSession(
                date=date(2026, 6, 10),
                day_number=1,
                completed_sets=[completed_set],
            ),
            CompletedSession(
                date=date(2026, 6, 11),
                day_number=4,
                completed_sets=[completed_set],
            ),
        ]

        self.assertEqual(suggest_next_training_day(history), 1)


class DoubleProgressionTest(unittest.TestCase):
    def test_mvp_load_increment_is_one_kg(self):
        self.assertEqual(MVP_LOAD_INCREMENT_KG, 1.0)

    def test_no_history_leaves_target_load_empty(self):
        plan = build_session_for_day(1, date(2026, 6, 11), CheckIn(raw_text=""))

        progressed = apply_double_progression(plan.planned_exercises[0], [])

        self.assertIsNone(progressed.prescribed_sets[0].target_load)

    def test_hitting_top_of_range_on_all_sets_adds_one_kg_and_resets_reps(self):
        previous_session = CompletedSession(
            date=date(2026, 6, 10),
            day_number=1,
            completed_sets=[
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=1,
                    actual_reps=15,
                    actual_load=20,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=2,
                    actual_reps=15,
                    actual_load=20,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=3,
                    actual_reps=15,
                    actual_load=20,
                ),
            ],
        )

        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text=""),
            completed_sessions=[previous_session],
        )

        target_loads = [
            prescribed_set.target_load
            for prescribed_set in plan.planned_exercises[0].prescribed_sets
        ]
        self.assertEqual(target_loads, [21, 21, 21])
        target_reps = [
            prescribed_set.target_reps
            for prescribed_set in plan.planned_exercises[0].prescribed_sets
        ]
        self.assertEqual(target_reps, [11, 11, 11])
        self.assertIn("add 1 kg", plan.planned_exercises[0].notes)

    def test_missing_top_of_range_repeats_latest_load_and_adds_two_reps(self):
        previous_session = CompletedSession(
            date=date(2026, 6, 10),
            day_number=1,
            completed_sets=[
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=1,
                    actual_reps=13,
                    actual_load=20,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=2,
                    actual_reps=12,
                    actual_load=20,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=3,
                    actual_reps=13,
                    actual_load=20,
                ),
            ],
        )

        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text=""),
            completed_sessions=[previous_session],
        )

        target_loads = [
            prescribed_set.target_load
            for prescribed_set in plan.planned_exercises[0].prescribed_sets
        ]
        self.assertEqual(target_loads, [20, 20, 20])
        target_reps = [
            prescribed_set.target_reps
            for prescribed_set in plan.planned_exercises[0].prescribed_sets
        ]
        self.assertEqual(target_reps, [15, 14, 15])
        self.assertIn("repeat latest load", plan.planned_exercises[0].notes)

    def test_progression_uses_latest_matching_exercise_exposure(self):
        older_session = CompletedSession(
            date=date(2026, 6, 9),
            day_number=1,
            completed_sets=[
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=1,
                    actual_reps=13,
                    actual_load=20,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=2,
                    actual_reps=13,
                    actual_load=20,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=3,
                    actual_reps=13,
                    actual_load=20,
                ),
            ],
        )
        latest_session = CompletedSession(
            date=date(2026, 6, 10),
            day_number=1,
            completed_sets=[
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=1,
                    actual_reps=12,
                    actual_load=22,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=2,
                    actual_reps=12,
                    actual_load=22,
                ),
                CompletedSet(
                    exercise_id="dumbbell-pullover",
                    set_number=3,
                    actual_reps=12,
                    actual_load=22,
                ),
            ],
        )

        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text=""),
            completed_sessions=[older_session, latest_session],
        )

        target_loads = [
            prescribed_set.target_load
            for prescribed_set in plan.planned_exercises[0].prescribed_sets
        ]
        self.assertEqual(target_loads, [22, 22, 22])


class PainFilterTest(unittest.TestCase):
    def test_exercise_muscle_map_tags_incline_bench_with_triceps(self):
        self.assertIn(
            Muscle.TRICEPS_BRACHII,
            EXERCISE_MUSCLE_MAP["barbell-incline-bench-press"],
        )

    def test_triceps_pain_removes_pressing_and_triceps_exercises(self):
        check_in = CheckIn(
            raw_text="My tricep hurts.",
            pain_or_injury="yes",
            pain_issues=[
                PainIssue(
                    affected_muscle=Muscle.TRICEPS_BRACHII,
                    severity="moderate",
                    notes="Tricep hurts.",
                )
            ],
        )

        plan = build_session_for_day(1, date(2026, 6, 11), check_in)
        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]

        self.assertNotIn("barbell-incline-bench-press", exercise_ids)
        self.assertNotIn("dumbbell-pullover", exercise_ids)
        self.assertIn("dumbbell-row", exercise_ids)
        self.assertIn("Pain filter removed", plan.notes)
        self.assertEqual([exercise.order for exercise in plan.planned_exercises], [1, 2, 3])

    def test_biceps_pain_removes_bicep_curl_exercises(self):
        check_in = CheckIn(
            raw_text="My bicep tore.",
            pain_or_injury="yes",
            pain_issues=[
                PainIssue(
                    affected_muscle=Muscle.BICEPS_BRACHII,
                    severity="severe",
                    notes="Bicep tore.",
                )
            ],
        )

        plan = build_session_for_day(4, date(2026, 6, 11), check_in)
        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]

        self.assertNotIn("ez-bar-biceps-curl", exercise_ids)
        self.assertNotIn("cross-body-hammer-curl", exercise_ids)
        self.assertIn("dumbbell-lateral-raise", exercise_ids)
        self.assertIn("Pain filter removed", plan.notes)

    def test_manual_soreness_constraints_bicep_tore_filters_bicep_curls(self):
        check_in = CheckIn(
            raw_text="",
            soreness="bicep tore",
            pain_or_injury="unsure",
        )

        plan = build_session_for_day(4, date(2026, 6, 11), check_in)
        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]

        self.assertNotIn("ez-bar-biceps-curl", exercise_ids)
        self.assertNotIn("cross-body-hammer-curl", exercise_ids)
        self.assertIn("Pain filter removed", plan.notes)

    def test_manual_soreness_constraints_calves_pain_filters_calf_raises(self):
        check_in = CheckIn(
            raw_text="",
            soreness="calves are cooked and in pain",
            pain_or_injury="unsure",
        )

        plan = build_session_for_day(4, date(2026, 6, 11), check_in)
        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]

        self.assertNotIn("barbell-standing-calf-raise", exercise_ids)
        self.assertIn("ez-bar-biceps-curl", exercise_ids)
        self.assertIn("Pain filter removed", plan.notes)

    def test_manual_pain_yes_with_bicep_text_filters_bicep_curls(self):
        check_in = CheckIn(
            raw_text="",
            soreness="bicep feels bad",
            pain_or_injury="yes",
        )

        plan = build_session_for_day(4, date(2026, 6, 11), check_in)
        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]

        self.assertNotIn("ez-bar-biceps-curl", exercise_ids)
        self.assertNotIn("cross-body-hammer-curl", exercise_ids)

    def test_unclear_or_unsure_pain_does_not_filter_until_clarified(self):
        check_in = CheckIn(
            raw_text="My arm feels weird.",
            pain_or_injury="unsure",
            pain_issues=[
                PainIssue(
                    affected_muscle=Muscle.TRICEPS_BRACHII,
                    severity="unsure",
                    notes="Arm feels weird.",
                )
            ],
        )
        plan = build_session_for_day(1, date(2026, 6, 11), check_in)

        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]
        self.assertIn("barbell-incline-bench-press", exercise_ids)

    def test_apply_pain_filter_returns_original_plan_when_no_known_pain_muscle(self):
        plan = build_session_for_day(1, date(2026, 6, 11), CheckIn(raw_text=""))

        filtered, notes = apply_pain_filter(plan.planned_exercises, CheckIn(raw_text=""))

        self.assertEqual(len(filtered), len(plan.planned_exercises))
        self.assertEqual(notes, [])


class TimeCompressionTest(unittest.TestCase):
    def test_sixty_minutes_keeps_full_plan(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text="60 min", time_available_minutes=60),
        )

        self.assertEqual(len(plan.planned_exercises), 5)
        self.assertEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            16,
        )
        self.assertNotIn("Time compression", plan.notes)

    def test_fifty_minutes_cuts_to_about_seventy_five_percent_sets(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text="50 min", time_available_minutes=50),
        )

        set_counts = [len(exercise.prescribed_sets) for exercise in plan.planned_exercises]
        self.assertEqual(sum(set_counts), 12)
        self.assertEqual(set_counts, [3, 3, 2, 2, 2])
        self.assertIn("Time compression", plan.notes)

    def test_thirty_five_minutes_cuts_to_about_sixty_percent_sets(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text="35 min", time_available_minutes=35),
        )

        set_counts = [len(exercise.prescribed_sets) for exercise in plan.planned_exercises]
        self.assertEqual(sum(set_counts), 10)
        self.assertEqual(set_counts[0], 3)
        self.assertEqual(set_counts[-1], 1)
        self.assertIn("35 minutes", plan.notes)

    def test_under_thirty_minutes_keeps_first_four_exercises_and_caps_sets(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(raw_text="25 min", time_available_minutes=25),
        )

        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]
        self.assertEqual(
            exercise_ids,
            [
                "dumbbell-pullover",
                "dumbbell-row",
                "barbell-incline-bench-press",
                "dumbbell-incline-chest-fly",
            ],
        )
        self.assertEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            8,
        )
        self.assertIn("under 30 minutes", plan.notes)

    def test_time_compression_runs_after_pain_filter(self):
        check_in = CheckIn(
            raw_text="25 min, tricep hurts",
            time_available_minutes=25,
            pain_or_injury="yes",
            pain_issues=[
                PainIssue(
                    affected_muscle=Muscle.TRICEPS_BRACHII,
                    severity="moderate",
                    notes="Tricep hurts.",
                )
            ],
        )

        plan = build_session_for_day(1, date(2026, 6, 11), check_in)
        exercise_ids = [exercise.exercise_id for exercise in plan.planned_exercises]

        self.assertNotIn("dumbbell-pullover", exercise_ids)
        self.assertNotIn("barbell-incline-bench-press", exercise_ids)
        self.assertLessEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            8,
        )
        self.assertIn("Pain filter removed", plan.notes)
        self.assertIn("Time compression", plan.notes)

    def test_apply_time_compression_can_be_called_directly(self):
        plan = build_session_for_day(1, date(2026, 6, 11), CheckIn(raw_text=""))

        compressed, notes = apply_time_compression(
            plan.planned_exercises,
            CheckIn(raw_text="45 min", time_available_minutes=45),
        )

        self.assertLess(
            sum(len(exercise.prescribed_sets) for exercise in compressed),
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
        )
        self.assertTrue(notes)


class ReadinessModifierTest(unittest.TestCase):
    def test_readiness_score_uses_neutral_defaults(self):
        score = readiness_score(CheckIn(raw_text=""))

        self.assertEqual(score, 3.0)

    def test_normal_readiness_sets_base_target_rir_without_cutting_sets(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(
                raw_text="60 min, okay sleep, medium energy",
                time_available_minutes=60,
                sleep_quality="okay",
                energy_level="medium",
                mood_stress="neutral",
            ),
        )

        self.assertEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            16,
        )
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_rir, 2)
        self.assertIn("Readiness: normal", plan.notes)

    def test_low_readiness_reduces_sets_and_adds_one_rir(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(
                raw_text="60 min, okay sleep, low energy",
                time_available_minutes=60,
                sleep_quality="okay",
                sleep_hours=6.5,
                energy_level="low",
                mood_stress="neutral",
            ),
        )

        self.assertEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            13,
        )
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_rir, 3)
        self.assertIn("Readiness: low", plan.notes)

    def test_very_low_readiness_reduces_sets_and_adds_two_rir(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(
                raw_text="60 min, poor sleep, low energy, sore and stressed",
                time_available_minutes=60,
                sleep_quality="poor",
                sleep_hours=4.5,
                energy_level="low",
                soreness="very sore",
                mood_stress="stressed",
            ),
        )

        self.assertEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            8,
        )
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_rir, 4)
        self.assertIn("Readiness: very low", plan.notes)
        self.assertIn("Sleep override", plan.notes)

    def test_high_readiness_does_not_auto_add_work_for_mvp(self):
        plan = build_session_for_day(
            1,
            date(2026, 6, 11),
            CheckIn(
                raw_text="90 min, good sleep, high energy, ready",
                time_available_minutes=90,
                sleep_quality="good",
                sleep_hours=8,
                energy_level="high",
                soreness="no soreness",
                mood_stress="ready",
            ),
        )

        self.assertEqual(
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
            16,
        )
        self.assertEqual(plan.planned_exercises[0].prescribed_sets[0].target_rir, 2)
        self.assertIn("Readiness: high", plan.notes)

    def test_apply_readiness_modifier_can_be_called_directly(self):
        plan = build_session_for_day(1, date(2026, 6, 11), CheckIn(raw_text=""))

        modified, notes = apply_readiness_modifier(
            plan.planned_exercises,
            CheckIn(
                raw_text="low energy",
                sleep_quality="okay",
                sleep_hours=6.5,
                energy_level="low",
                mood_stress="neutral",
            ),
        )

        self.assertLess(
            sum(len(exercise.prescribed_sets) for exercise in modified),
            sum(len(exercise.prescribed_sets) for exercise in plan.planned_exercises),
        )
        self.assertTrue(notes)


if __name__ == "__main__":
    unittest.main()