module Stimulus = struct
  type extension = Forced_reps | Negatives | Rest_pause | Static_hold

  type outcome =
    | Positive_failure
    | Beyond_failure of extension * extension list

  let extensions_of_outcome = function
    | Positive_failure -> []
    | Beyond_failure (first, rest) -> first :: rest

  let pp_extension ppf e =
    Format.pp_print_string ppf
      (match e with
      | Forced_reps -> "forced reps"
      | Negatives -> "negatives"
      | Rest_pause -> "rest-pause"
      | Static_hold -> "static hold")

  module Effort = struct
    type t = {
      exercise : Exercise.t;
      load : float;
      reps : int;
      outcome : outcome;
    }

    let make ~exercise ~load ~reps ~outcome = { exercise; load; reps; outcome }
    let exercise t = t.exercise
    let load t = t.load
    let reps t = t.reps
    let outcome t = t.outcome
    let extensions t = extensions_of_outcome t.outcome

    let pp ppf t =
      Format.fprintf ppf "%a %g kg x %d" Exercise.pp t.exercise t.load t.reps;
      match extensions t with
      | [] -> ()
      | es ->
          Format.fprintf ppf " (%a)"
            (Format.pp_print_list
               ~pp_sep:(fun ppf () -> Format.pp_print_string ppf ", ")
               pp_extension)
            es
  end

  type delivery =
    | Single of Effort.t
    | Pair of { first : Effort.t; second : Effort.t }

  type t = { delivery : delivery }

  let make delivery = { delivery }
  let delivery t = t.delivery

  let efforts t =
    match t.delivery with
    | Single effort -> [ effort ]
    | Pair { first; second } -> [ first; second ]

  let exercises t = List.map Effort.exercise (efforts t)
  let extensions t = List.concat_map Effort.extensions (efforts t)
  let is_extended t = extensions t <> []

  let pp ppf t =
    match t.delivery with
    | Single effort -> Effort.pp ppf effort
    | Pair { first; second } ->
        Format.fprintf ppf "%a then %a" Effort.pp first Effort.pp second
end

module Feedback = struct
  type level = Very_poor | Poor | Fair | Good | Very_good

  let level_to_score = function
    | Very_poor -> 1
    | Poor -> 2
    | Fair -> 3
    | Good -> 4
    | Very_good -> 5

  let level_of_score = function
    | 1 -> Some Very_poor
    | 2 -> Some Poor
    | 3 -> Some Fair
    | 4 -> Some Good
    | 5 -> Some Very_good
    | _ -> None

  type signal =
    | Sleep of level
    | Appetite of level
    | Readiness of level
    | Motivation of level
    | Difficulty of level
    | Pain
    | Injury
    | Preparation_insufficient

  type t = { reported_at : Recovery.timestamp; signals : signal list }
  type error = Duplicate_signal of signal

  exception Invalid of error

  let same_category left right =
    match (left, right) with
    | Sleep _, Sleep _
    | Appetite _, Appetite _
    | Readiness _, Readiness _
    | Motivation _, Motivation _
    | Difficulty _, Difficulty _
    | Pain, Pain
    | Injury, Injury
    | Preparation_insufficient, Preparation_insufficient ->
        true
    | _ -> false

  let make ~reported_at signals =
    let rec validate seen = function
      | [] -> { reported_at; signals }
      | signal :: rest ->
          if List.exists (same_category signal) seen then
            raise (Invalid (Duplicate_signal signal))
          else validate (signal :: seen) rest
    in
    validate [] signals

  let reported_at t = t.reported_at
  let signals t = t.signals
end

module Workout = struct
  type shape = As_single | As_pair
  type completeness = Complete | Incomplete

  type error =
    | Not_prescribed of Exercise.id
    | Delivery_mismatch of {
        exercise : Exercise.id;
        prescribed : shape;
        logged : shape;
      }
    | No_such_slot of int

  exception Invalid of error

  type t = {
    prescription : Prescription.Workout.t;
    clearance : Recovery.clearance;
    started_at : Recovery.timestamp;
    ended_at : Recovery.timestamp option;
    performed : (int * Stimulus.t) list;
  }

  let pp_shape ppf = function
    | As_single -> Format.pp_print_string ppf "a single set"
    | As_pair -> Format.pp_print_string ppf "a pair"

  let pp_error ppf = function
    | Not_prescribed id ->
        Format.fprintf ppf "%s is not prescribed for this workout"
          (id :> string)
    | Delivery_mismatch { exercise; prescribed; logged } ->
        Format.fprintf ppf "%s is prescribed as %a but was logged as %a"
          (exercise :> string)
          pp_shape prescribed pp_shape logged
    | No_such_slot slot ->
        Format.fprintf ppf "slot %d is not part of this workout" slot

  let start prescription ~clearance ~started_at =
    { prescription; clearance; started_at; ended_at = None; performed = [] }

  let prescription t = t.prescription
  let clearance t = t.clearance
  let started_at t = t.started_at
  let ended_at t = t.ended_at
  let is_finished t = Option.is_some t.ended_at
  let stimuli t = List.map snd (List.rev t.performed)
  let performed t = List.rev t.performed

  let duration t =
    Option.map
      (fun ended -> Recovery.elapsed ~since:t.started_at ~now:ended)
      t.ended_at

  let prescribed_shape p =
    match Prescription.Stimulus.delivery p with
    | Prescription.Stimulus.Single _ -> As_single
    | Prescription.Stimulus.Pre_exhaust _ -> As_pair

  let logged_shape s =
    match Stimulus.delivery s with
    | Stimulus.Single _ -> As_single
    | Stimulus.Pair _ -> As_pair

  let fills ~prescribed ~logged p =
    Exercise.equal prescribed logged
    || List.exists (Exercise.equal logged)
         (Prescription.Stimulus.allowed_substitutes p)
       && Exercise.may_substitute ~original:prescribed ~candidate:logged

  let mentions p s =
    List.for_all
      (fun e -> Prescription.Stimulus.permits p e)
      (Stimulus.exercises s)

  let conforms p s =
    let logged = Stimulus.exercises s in
    let prescribed = Prescription.Stimulus.exercises p in
    List.length logged = List.length prescribed
    && List.for_all2
         (fun prescribed logged -> fills ~prescribed ~logged p)
         prescribed logged

  let indexed t =
    List.mapi (fun i p -> (i, p)) (Prescription.Workout.stimuli t.prescription)

  let answered t = List.map fst t.performed

  let outstanding t =
    indexed t |> List.filter (fun (i, _) -> not (List.mem i (answered t)))

  let unperformed t = List.map snd (outstanding t)
  let completeness t = if outstanding t = [] then Complete else Incomplete

  (* Distinct answered slots, so recorded extra volume never inflates the
     count past the prescription length. *)
  let filled_slots t = List.length (List.sort_uniq Int.compare (answered t))

  let add_stimulus t s =
    let candidates = List.filter (fun (_, p) -> mentions p s) (indexed t) in
    let matching = List.filter (fun (_, p) -> conforms p s) candidates in
    let leading = Exercise.id (List.hd (Stimulus.exercises s)) in
    match (candidates, matching) with
    | [], _ -> raise (Invalid (Not_prescribed leading))
    | (_, p) :: _, [] ->
        raise
          (Invalid
             (Delivery_mismatch
                {
                  exercise = leading;
                  prescribed = prescribed_shape p;
                  logged = logged_shape s;
                }))
    | _, matching ->
        let unanswered =
          List.filter (fun (i, _) -> not (List.mem i (answered t))) matching
        in
        let i, _ =
          match unanswered with chosen :: _ -> chosen | [] -> List.hd matching
        in
        { t with performed = (i, s) :: t.performed }

  (* Validate a stimulus against one named slot, raising as [add_stimulus]
     does. Shared by the slot-targeted operations. *)
  let prescription_at t slot =
    match List.nth_opt (Prescription.Workout.stimuli t.prescription) slot with
    | Some p -> p
    | None -> raise (Invalid (No_such_slot slot))

  let check_against p s =
    if not (mentions p s) then
      raise
        (Invalid (Not_prescribed (Exercise.id (List.hd (Stimulus.exercises s)))))
    else if not (conforms p s) then
      raise
        (Invalid
           (Delivery_mismatch
              {
                exercise = Exercise.id (List.hd (Stimulus.exercises s));
                prescribed = prescribed_shape p;
                logged = logged_shape s;
              }))

  let record_at t ~slot s =
    check_against (prescription_at t slot) s;
    { t with performed = (slot, s) :: t.performed }

  let replace_stimulus t ~slot s =
    check_against (prescription_at t slot) s;
    (* A correction targets the slot; it never adds to the recorded volume.
       Replace the slot's fill in place, keeping its position in performance
       order. When the slot was empty, record it as its first fill. *)
    if List.mem_assoc slot t.performed then
      let replaced = ref false in
      let performed =
        List.map
          (fun (i, existing) ->
            if i = slot && not !replaced then begin
              replaced := true;
              (i, s)
            end
            else (i, existing))
          t.performed
      in
      { t with performed }
    else { t with performed = (slot, s) :: t.performed }

  let finish t ~ended_at =
    match t.ended_at with
    | None -> { t with ended_at = Some ended_at }
    | Some _ -> t

  let pp ppf t =
    Format.fprintf ppf "%a (%d of %d)" Prescription.Workout.pp t.prescription
      (List.length t.performed)
      (List.length (Prescription.Workout.stimuli t.prescription))
end

module Log = struct
  type t = Workout.t list

  type observation = {
    exercise : Exercise.t;
    effort : Stimulus.Effort.t;
    performed_at : Recovery.timestamp;
  }

  let empty = []
  let add t workout = workout :: t

  let started workout =
    Recovery.timestamp_to_unix_seconds (Workout.started_at workout)

  (* Sorted on read, so workouts need not be added in order. *)
  let chronological t =
    List.sort (fun a b -> Int.compare (started a) (started b)) t

  let workouts t = List.rev (chronological t)

  let last_prescription t =
    match workouts t with
    | [] -> None
    | latest :: _ -> Some (Workout.prescription latest)

  let observations_of workout =
    Workout.stimuli workout
    |> List.concat_map Stimulus.efforts
    |> List.map (fun effort ->
        {
          exercise = Stimulus.Effort.exercise effort;
          effort;
          performed_at = Workout.started_at workout;
        })

  let observations t exercise =
    chronological t
    |> List.concat_map observations_of
    |> List.filter (fun o -> Exercise.equal o.exercise exercise)

  let readiness t ~now ~recommended =
    let finished =
      chronological t
      |> List.filter_map (fun w -> Workout.ended_at w)
      |> List.rev
    in
    match finished with
    | [] -> Recovery.Ready
    | last :: _ ->
        Recovery.evaluate_readiness
          ~elapsed:(Recovery.elapsed ~since:last ~now)
          ~recommended

  let pp ppf t = Format.fprintf ppf "%d workouts" (List.length t)
end