(* Common util. Should probably put this in a module or something.. *)
let get_one_line file =
  try Some (input_line file) with
    End_of_file -> None

let get_lines file =
  let rec input lines =
    match get_one_line file with
      Some line -> input (line :: lines)
    | None -> List.rev lines
  in input []

let read_file file =
  let channel = open_in(file) in
  get_lines channel

let (>>) f g x = g(f(x))

let list_max l = List.fold_left (fun acc x -> if x > acc then x else acc) (List.hd l) l

let map_tuple f (a, b) = (f a, f b)

let memo_rec f =
  let h = Hashtbl.create 1000 in
  let rec g x =
    try Hashtbl.find h x
    with Not_found ->
      let y = f g x in
      Hashtbl.add h x y;
      y
  in g

(* --- *)

let rec collect_diffs = function
  | x :: y :: [] -> [y - x] (* 2 elemnts left *)
  | x :: y :: z -> y - x :: collect_diffs (y :: z)

let count_steps = (* returns a tuple of counts of 1's and 3's *)
  List.partition (Int.equal 1)
  >> map_tuple List.length

let count_valid self = function
  | x :: y :: [] -> if y - x < 4 then 1 else 0
  | x :: y :: rest ->
      if y - x > 3
      then 0
      else self (x :: rest) + self (y :: rest)

let count_valid_memo = memo_rec count_valid

let () =
  let lines = read_file "day10.txt"
    |> List.map int_of_string
  in

  let adapter_jolts = 3 + list_max lines in

  let sorted = lines
    |> List.cons adapter_jolts (* add our adapter *)
    |> List.cons 0 (* add the outlet *)
    |> List.sort (fun x y -> if x > y then 1 else if x < y then -1 else 0)
  in

  sorted
    |> collect_diffs
    |> count_steps
    |> fun (x, y) -> x * y
    |> Printf.printf "prod: %i\n";

  sorted
    |> count_valid_memo
    |> Printf.printf "\nsum: %i\n"