module Effort = Evidence.Stimulus.Effort

let beats ~previous ~current =
  let load = Float.compare (Effort.load current) (Effort.load previous) in
  let reps = Int.compare (Effort.reps current) (Effort.reps previous) in
  load > 0 || (load = 0 && reps > 0)

(* The mirror of [beats]: [current] is strictly worse than [previous] — less
   load, or equal load and fewer reps. The mark of decompensation. *)
let declines ~previous ~current =
  let load = Float.compare (Effort.load current) (Effort.load previous) in
  let reps = Int.compare (Effort.reps current) (Effort.reps previous) in
  load < 0 || (load = 0 && reps < 0)

type assessment = Progressing | Stalled | Decompensating
type error = Insufficient_data

let pp_assessment ppf = function
  | Progressing -> Format.pp_print_string ppf "progressing"
  | Stalled -> Format.pp_print_string ppf "stalled"
  | Decompensating -> Format.pp_print_string ppf "decompensating"

let pp_error ppf Insufficient_data =
  Format.pp_print_string ppf "not enough record to judge"

let stall_window = Recovery.days 14
let rec last = function [] -> None | [ x ] -> Some x | _ :: rest -> last rest

(* The most recent observation that improved on the one before it. *)
let last_advance observations =
  let rec scan found = function
    | (previous : Evidence.Observation.t) :: (current :: _ as rest) ->
        let found =
          if
            beats
              ~previous:(Evidence.Observation.effort previous)
              ~current:(Evidence.Observation.effort current)
          then Some current
          else found
        in
        scan found rest
    | _ -> found
  in
  scan None observations

let elapsed_between (a : Evidence.Observation.t) (b : Evidence.Observation.t) =
  Recovery.duration_to_seconds
    (Recovery.elapsed
       ~since:(Evidence.Observation.performed_at a)
       ~now:(Evidence.Observation.performed_at b))

let assess observations =
  let window = Recovery.duration_to_seconds stall_window in
  match (observations, last observations) with
  | ([] | [ _ ]), _ | _, None -> Error Insufficient_data
  | first :: _, Some latest -> (
      (* After the window passes without progress, separate a halt from a loss.
         Strength below the reference is decompensation. Strength that holds at
         the reference is a stall. The reference is the peak reached (the last
         advance), or the first observation when the record never advanced. *)
      let effort = Evidence.Observation.effort in
      let verdict ~reference =
        if declines ~previous:(effort reference) ~current:(effort latest) then
          Decompensating
        else Stalled
      in
      match last_advance observations with
      | Some advance ->
          if elapsed_between advance latest >= window then
            Ok (verdict ~reference:advance)
          else Ok Progressing
      | None ->
          (* Never advanced. Only a stall once the record is long enough to say
             so. Otherwise two sessions a day apart would condemn the routine. *)
          if elapsed_between first latest >= window then
            Ok (verdict ~reference:first)
          else Error Insufficient_data)

type remedy =
  | Lay_off_then_reduce of {
      lay_off : Recovery.duration;
      drop_stimuli_per_workout : int;
      extra_rest : Recovery.duration;
    }

let remedy = function
  | Progressing -> None
  | Stalled | Decompensating ->
      Some
        (Lay_off_then_reduce
           {
             lay_off = Recovery.days 7;
             drop_stimuli_per_workout = 1;
             extra_rest = Recovery.days 1;
           })

let pp_remedy ppf
    (Lay_off_then_reduce { lay_off; drop_stimuli_per_workout; extra_rest }) =
  Format.fprintf ppf
    "take %a off, then drop %d stimulus per workout and add %a between workouts"
    Recovery.pp_duration lay_off drop_stimuli_per_workout Recovery.pp_duration
    extra_rest

let load_increase_trigger = 12
let load_increase ~current = (current *. 1.10, current *. 1.20)

type load_verdict = Hold | Increase of float * float | Too_heavy

(* The trigger is absolute. The window ceiling never gates the verdict. Only
   the window floor does. *)
let judge_load ~rep_range effort =
  let reps = Effort.reps effort in
  let min_reps, _ = Prescription.Rep_range.bounds rep_range in
  if reps >= load_increase_trigger then
    let low, high = load_increase ~current:(Effort.load effort) in
    Increase (low, high)
  else if reps < min_reps then Too_heavy
  else Hold

let pp_load_verdict ppf = function
  | Hold -> Format.pp_print_string ppf "hold the load"
  | Increase (low, high) ->
      Format.fprintf ppf "raise the load to %g-%g kg" low high
  | Too_heavy -> Format.pp_print_string ppf "load is too heavy for the window"

type diagnostic =
  | Excess_volume of int
  | Extensions_on_every_stimulus of int
  | Trained_under_recovered of int

let excess_volume workout =
  List.length (Evidence.Workout.stimuli workout)
  > Evidence.Workout.filled_slots workout

let all_extended workout =
  match Evidence.Workout.stimuli workout with
  | [] -> false
  | stimuli -> List.for_all Evidence.Stimulus.is_extended stimuli

let under_recovered workout =
  match Recovery.basis (Evidence.Workout.clearance workout) with
  | Recovery.Recovered -> false
  | Recovery.Overridden _ -> true

let diagnose log =
  let workouts = Evidence.Log.workouts log in
  let count predicate = List.length (List.filter predicate workouts) in
  let excessive = count excess_volume in
  let extended = count all_extended in
  let overridden = count under_recovered in
  let when_present n d = if n > 0 then [ d n ] else [] in
  when_present excessive (fun n -> Excess_volume n)
  @ when_present extended (fun n -> Extensions_on_every_stimulus n)
  @ when_present overridden (fun n -> Trained_under_recovered n)

let pp_diagnostic ppf = function
  | Excess_volume n ->
      Format.fprintf ppf "%d workout(s) recorded more stimuli than prescribed" n
  | Extensions_on_every_stimulus n ->
      Format.fprintf ppf "%d workout(s) extended every stimulus beyond failure"
        n
  | Trained_under_recovered n ->
      Format.fprintf ppf "%d workout(s) begun before recovery was complete" n

type workout_comparison = {
  exercise : Exercise.t;
  previous : Evidence.Stimulus.Effort.t;
  current : Evidence.Stimulus.Effort.t;
}

(* The latest observation strictly before [at]. This prevents a correction from
   comparing an exercise against itself when the recorded observation shares its
   workout's start time. Callers scope the observation series by exercise before
   they call this helper. *)
let latest_before ~at observations =
  observations
  |> List.filter (fun (observation : Evidence.Observation.t) ->
      Recovery.duration_to_seconds
        (Recovery.elapsed
           ~since:(Evidence.Observation.performed_at observation)
           ~now:at)
      > 0)
  |> List.sort (fun a b ->
      Int.compare
        (Recovery.timestamp_to_unix_seconds
           (Evidence.Observation.performed_at a))
        (Recovery.timestamp_to_unix_seconds
           (Evidence.Observation.performed_at b)))
  |> last

(* The latest observation strictly before the workout started. This prevents a
   correction of the saved workout from comparing an exercise against itself. *)
let previous_effort log ~started exercise =
  Evidence.Log.observations log exercise
  |> latest_before ~at:started
  |> Option.map Evidence.Observation.effort

let workout_comparisons log workout =
  let started = Evidence.Workout.started_at workout in
  Evidence.Workout.stimuli workout
  |> List.concat_map Evidence.Stimulus.efforts
  |> List.fold_left
       (fun (seen, comparisons) current ->
         let exercise = Evidence.Stimulus.Effort.exercise current in
         if List.exists (Exercise.equal exercise) seen then (seen, comparisons)
         else
           let comparison =
             Option.map
               (fun previous -> { exercise; previous; current })
               (previous_effort log ~started exercise)
           in
           (exercise :: seen, Option.to_list comparison @ comparisons))
       ([], [])
  |> snd |> List.rev

let workout_changes log workout =
  workout_comparisons log workout
  |> List.filter (fun comparison ->
      beats ~previous:comparison.previous ~current:comparison.current)

(* How many distinct exercises in [workout] beat their most recent earlier
   observation in [log]. Progress is the Heavy Duty signal: an increase in reps
   or load. This module judges it, so a caller receives a count and never
   re-derives the rule. An exercise with no earlier observation does not
   count. *)
let workout_progress log workout = List.length (workout_changes log workout)

(* The most recent observation of [target]'s exercise strictly before [target]
   was performed, taken from [observations] (any order). This applies the same
   baseline rule as [previous_effort], but over a bare observation series rather
   than a log. Imported evidence is not a native workout. A caller can still
   compare it against the observations before it. *)
let previous_observation observations (target : Evidence.Observation.t) =
  let exercise = Evidence.Observation.exercise target in
  let performed = Evidence.Observation.performed_at target in
  observations
  |> List.filter (fun (o : Evidence.Observation.t) ->
      Exercise.equal (Evidence.Observation.exercise o) exercise)
  |> latest_before ~at:performed