let pin_count = 5
let height = 7

let read_lock_or_key lst =
  let result = Array.make pin_count 0 in
  let add_node i c = if c = '#' then result.(i) <- result.(i) + 1 in
  List.iter (String.iteri add_node) lst;
  result |> Array.to_list

let locks_and_keys_of_list =
  let rec impl locks keys = function
    | [] -> (locks, keys)
    | "" :: t -> impl locks keys t
    | a :: b :: c :: d :: e :: f :: g :: t ->
        let h = read_lock_or_key [ a; b; c; d; e; f; g ] in
        if a = String.make pin_count '#' then impl locks (h :: keys) t
        else impl (h :: locks) keys t
    | _ -> failwith "locks_and_keys_of_list"
  in
  impl [] []

let locks_and_keys_of_file fname =
  Aoc.strings_of_file fname |> locks_and_keys_of_list

let lock_key_fit lock key =
  List.map2 ( + ) lock key |> List.for_all (( >= ) height)

let count_keys keys lock = List.filter (lock_key_fit lock) keys |> List.length

let count_locks_and_keys (locks, keys) =
  List.map (count_keys keys) locks |> List.fold_left ( + ) 0

let _ =
  Aoc.main locks_and_keys_of_file [ (string_of_int, count_locks_and_keys) ]