(*
 * SPDX-FileCopyrightText: Copyright 2025 Alexandre Jesus <https://adbjesus.com>
 *
 * SPDX-License-Identifier: GPL-3.0-or-later
 *)

let parse_data ch =
  In_channel.input_all ch |> String.trim |> String.split_on_char '\n'

(** Computes the maximum joltage using [n] banks in [O(|bat|)] time.

    The algorithm works as follows:
    - Keep a stack of the accepted banks, initially the stack is empty
    - Keep track of the number [k = |bat| - n] of banks that can be discarded

    Then process banks sequentially according to the following steps:
    - While the stack is not empty, [k > 0], and the current bank is bigger
      than the one at the top of the stack, pop a bank from the stack and
      reduce [k] by 1
    - Insert the bank into the stack if the size of the stack is less then [n]
    - Otherwise, don't insert the bank, and reduce [k] by 1

    At the end the stack contains the relevant banks, in reverse order.

    Since each bank is inserted, and removed, at most once from the
    stack, and all operations are [O(1)] the complexity is
    [O(|bat|)].
 *)
let joltage n bat =
  let rec fn s k m l =
    match Seq.uncons l with
    | None ->
        let zero = int_of_char '0' in
        let digit c = int_of_char c - zero in
        List.rev_map digit s |> List.fold_left (fun a v -> (a * 10) + v) 0
    | Some (h, t) -> (
      match s with
      | sh :: st when k > 0 && sh < h ->
          fn st (k - 1) (m - 1) l
      | s when m < n ->
          fn (h :: s) k (m + 1) t
      | s ->
          fn s (k - 1) m t )
  in
  let k = String.length bat - n in
  fn [] k 0 (String.to_seq bat)

let solve ch n =
  parse_data ch
  |> List.map (joltage n)
  |> List.fold_left ( + ) 0
  |> Printf.printf "%d\n"

let part1 ch = solve ch 2

let part2 ch = solve ch 12