type verdict = Passed | Failed
type 'a prefix = { length : int; last : 'a }

type 'a step =
  | Probes of 'a prefix list
  | Culprit of { culprit : 'a prefix; passing : 'a prefix option }

(* [failing] is the shortest prefix observed to fail, the whole range at
   first; [passed] the prefixes shorter than it observed to pass, sorted and
   distinct. The pair is the same whatever the order its verdicts came in. *)
type 'a t = { suspects : 'a array; failing : int; passed : int list }

let err_no_probe ~range length =
  Fmt.kstr
    (fun m -> Error (`Msg m))
    "no probe of a range of %d ends at suspect %d" range length

let v = function
  | [] -> Error (`Msg "a red range holds at least one suspect")
  | suspects ->
      let suspects = Array.of_list suspects in
      Ok { suspects; failing = Array.length suspects; passed = [] }

let prefix t length = { length; last = t.suspects.(length - 1) }

(* The longest prefix observed to pass short of [failing]; 0, the empty
   prefix, when none is. *)
let passing t = List.fold_left max 0 t.passed

let passing_prefix t =
  match passing t with 0 -> None | low -> Some (prefix t low)

(* The suspects left are [passing + 1 .. failing], and [failing]'s prefix
   fails. The [j]-th of the [ways - 1] probes ends the [j]-th of [ways] parts
   whose sizes differ by at most one. *)
let next ?(ways = 2) t =
  if ways < 2 then invalid_arg "Bisect.next: ways must be at least 2";
  let low = passing t in
  if t.failing = low + 1 then
    Culprit { culprit = prefix t t.failing; passing = passing_prefix t }
  else
    let left = t.failing - low in
    let parts = min ways left in
    Probes
      (List.init (parts - 1) (fun j ->
           prefix t (low + ((j + 1) * left / parts))))

(* Each probe added halves the longest stretch of suspects left that no
   probe under way bounds, the earliest of equals. *)
let more ?(ways = 2) ~under_way t =
  if ways < 2 then invalid_arg "Bisect.more: ways must be at least 2";
  let low = passing t in
  let rec go cuts added n =
    let rec longest best = function
      | a :: (b :: _ as rest) ->
          let best =
            match best with
            | Some (x, y) when y - x >= b - a -> best
            | _ when b - a >= 2 -> Some (a, b)
            | _ -> best
          in
          longest best rest
      | [ _ ] | [] -> best
    in
    if n <= 0 then added
    else
      match longest None (List.sort_uniq Int.compare cuts) with
      | None -> added
      | Some (a, b) ->
          let m = a + ((b - a) / 2) in
          go (m :: cuts) (m :: added) (n - 1)
  in
  let inside l = l > low && l < t.failing in
  let under_way = List.filter inside under_way in
  go (low :: t.failing :: under_way) [] (ways - 1 - List.length under_way)
  |> List.sort Int.compare
  |> List.map (prefix t)

let observe t ~length verdict =
  let n = Array.length t.suspects in
  if length < 1 || length >= n then err_no_probe ~range:n length
  else
    match verdict with
    | Failed ->
        let failing = min t.failing length in
        Ok
          {
            t with
            failing;
            passed = List.filter (fun p -> p < failing) t.passed;
          }
    | Passed ->
        if length >= t.failing then Ok t
        else
          Ok { t with passed = List.sort_uniq Int.compare (length :: t.passed) }

let left t = t.failing - passing t

(* A round leaves at most the largest of its parts, [ceil (left / ways)]. *)
let rounds ?(ways = 2) n =
  if ways < 2 then invalid_arg "Bisect.rounds: ways must be at least 2";
  let rec go acc n =
    if n <= 1 then acc else go (acc + 1) ((n + ways - 1) / ways)
  in
  go 0 n

let control t =
  Fmt.str "of %d, failing %d, passed %a" (Array.length t.suspects) t.failing
    Fmt.(brackets (list ~sep:comma int))
    t.passed

let pp ppf t =
  Fmt.pf ppf "bisection of %d suspects, %d to %d left" (Array.length t.suspects)
    (passing t + 1)
    t.failing