tracker/progression.py

from __future__ import annotations

import math

from pydantic import BaseModel

from tracker.models import Exercise, ProgramExercise, WorkoutLog
from tracker.starting_load_research import (
    LB_TO_KG,
    equipment_has_dumbbell,
    resolve_research_starting_weight_kg,
)


class SuggestedSet(BaseModel):
    """What the next session should look like for one exercise."""

    weight_kg: float | None
    target_reps: int
    note: str
    load_reference: str | None = None


def quantize_half_kg(value: float) -> float:
    """Nearest 0.5 kg (typical plate resolution). Half-up tiebreak."""
    return math.floor(value * 2 + 0.5) / 2


def quantize_suggested_load_kg(
    value_kg: float,
    exercise: Exercise,
    dumbbell_quantum_lb: float | None,
) -> float:
    """Barbell-style half-kg, or dumbbell rack steps in whole lb (default from config)."""
    if dumbbell_quantum_lb is None or dumbbell_quantum_lb <= 0:
        return quantize_half_kg(value_kg)
    if not equipment_has_dumbbell(exercise.equipment):
        return quantize_half_kg(value_kg)
    lb = value_kg / LB_TO_KG
    quantum = dumbbell_quantum_lb
    rounded_lb = round(lb / quantum) * quantum
    if lb > 0:
        rounded_lb = max(quantum, rounded_lb)
    return rounded_lb * LB_TO_KG


def suggest_next(
    pe: ProgramExercise,
    last_session_logs: list[WorkoutLog],
    exercise: Exercise,
    body_weight_kg: float | None,
    *,
    calibration_bench_lb_per_hand: float | None = None,
    calibration_squat_lb_per_hand: float | None = None,
    dumbbell_weight_quantum_lb: float | None = None,
) -> SuggestedSet:
    """Compute suggested weight and reps from program rules and logged performance.

    **Double progression** (reps within range, then add load and reset reps) matches
    common prescription in ACSM progression models for healthy adults (Med Sci Sports
    Exerc 2009;41(3):687-708).

    **First session** (no logs): program ``starting_weight_kg`` overrides; else
    research-based starter from ``starting_load_research`` and latest body weight.

    **Later sessions** use logged loads from the last session for this slot.

    **Failure** (any set below ``target_rep_min``): reduce weight by one **program**
    ``weight_increment_kg`` step (your chosen bar/plate resolution), quantized to
    0.5 kg. If increment is zero (e.g. bodyweight-only slot), use a 0.5 kg minimum
    step when load is positive; at ~0 kg load, keep load at 0 and reset rep target
    to the bottom of the range.
    """
    if not last_session_logs:
        if pe.weight_increment_kg <= 0:
            return SuggestedSet(
                weight_kg=None,
                target_reps=pe.target_rep_min,
                note="No prior log: log reps for each set.",
            )
        if pe.starting_weight_kg is not None and pe.starting_weight_kg > 0:
            return SuggestedSet(
                weight_kg=quantize_suggested_load_kg(
                    pe.starting_weight_kg,
                    exercise,
                    dumbbell_weight_quantum_lb,
                ),
                target_reps=pe.target_rep_min,
                note="Starting load from program (PATCH exercise if your gym differs).",
                load_reference="Program exercise override.",
            )
        research_kg, ref = resolve_research_starting_weight_kg(
            exercise,
            body_weight_kg,
            calibration_bench_lb_per_hand=calibration_bench_lb_per_hand,
            calibration_squat_lb_per_hand=calibration_squat_lb_per_hand,
        )
        if research_kg is not None and research_kg > 0:
            return SuggestedSet(
                weight_kg=quantize_suggested_load_kg(
                    research_kg,
                    exercise,
                    dumbbell_weight_quantum_lb,
                ),
                target_reps=pe.target_rep_min,
                note="Evidence-based starter load for this movement; adjust to RPE and plates.",
                load_reference=ref,
            )
        if body_weight_kg is None or body_weight_kg <= 0:
            return SuggestedSet(
                weight_kg=None,
                target_reps=pe.target_rep_min,
                note=(
                    "Log body weight above to compute a research-based starter "
                    f"from your mass. ({ref})"
                ),
                load_reference=ref,
            )
        return SuggestedSet(
            weight_kg=None,
            target_reps=pe.target_rep_min,
            note=(
                "Enter the load you will use; no default fraction for this equipment. "
                f"({ref})"
            ),
            load_reference=ref,
        )

    last_weight = max(log.weight_kg for log in last_session_logs)
    all_hit_max = all(log.reps_done >= pe.target_rep_max for log in last_session_logs)
    any_below_min = any(log.reps_done < pe.target_rep_min for log in last_session_logs)

    if all_hit_max:
        if pe.weight_increment_kg <= 0:
            return SuggestedSet(
                weight_kg=None,
                target_reps=pe.target_rep_min,
                note="All sets at top of rep range: reset to bottom; progress with reps or harder variation.",
            )
        return SuggestedSet(
            weight_kg=quantize_suggested_load_kg(
                last_weight + pe.weight_increment_kg,
                exercise,
                dumbbell_weight_quantum_lb,
            ),
            target_reps=pe.target_rep_min,
            note="All sets at top of rep range: add load, reset reps.",
        )

    if any_below_min:
        if pe.weight_increment_kg <= 0:
            return SuggestedSet(
                weight_kg=None,
                target_reps=pe.target_rep_min,
                note="Below rep target last session: reset to bottom of rep range.",
            )
        step_down = pe.weight_increment_kg if pe.weight_increment_kg > 0 else 0.5
        if last_weight < 0.25:
            return SuggestedSet(
                weight_kg=0.0,
                target_reps=pe.target_rep_min,
                note="Below rep target: reset to bottom of rep range at bodyweight.",
            )
        reduced = quantize_suggested_load_kg(
            last_weight - step_down,
            exercise,
            dumbbell_weight_quantum_lb,
        )
        floor_kg = max(0.5, reduced)
        return SuggestedSet(
            weight_kg=floor_kg,
            target_reps=pe.target_rep_min,
            note="Below rep target last session: one increment lighter.",
        )

    best_reps = max(log.reps_done for log in last_session_logs)
    rep_note = (
        "Add reps within range before increasing load."
        if pe.weight_increment_kg > 0
        else "Add reps within range before a harder variation."
    )
    return SuggestedSet(
        weight_kg=(
            quantize_suggested_load_kg(last_weight, exercise, dumbbell_weight_quantum_lb)
            if pe.weight_increment_kg > 0
            else None
        ),
        target_reps=min(best_reps + 1, pe.target_rep_max),
        note=rep_note,
    )