(** [towels_of_strings lst] returns a pair containing a list of available towels
    and a list of patterns wanted. *)
let towels_of_strings = function
  | h :: "" :: t ->
      let re = Str.regexp "[, ]+" in
      let h = Str.split re h in
      (h, t)
  | _ -> failwith "towels_of_strings"

let towels_of_file fname = Aoc.strings_of_file fname |> towels_of_strings
let memo = Hashtbl.create 1000

let rec memoize_has_pattern pattern towels =
  let pattern_len = String.length pattern in
  let rec has_pattern = function
    | [] -> false
    | h :: t ->
        let towel_len = String.length h in
        if String.starts_with ~prefix:h pattern then
          memoize_has_pattern
            (String.sub pattern towel_len (pattern_len - towel_len))
            towels
          || has_pattern t
        else has_pattern t
  in
  match Hashtbl.find_opt memo pattern with
  | Some x -> x
  | None ->
      let x = if pattern_len = 0 then true else has_pattern towels in
      Hashtbl.add memo pattern x;
      x

let has_match towels pattern = memoize_has_pattern pattern towels
let memo2 = Hashtbl.create 1000

let rec memoize_count_matches pattern towels =
  let pattern_len = String.length pattern in
  let rec has_pattern = function
    | [] -> 0
    | h :: t ->
        let towel_len = String.length h in
        if String.starts_with ~prefix:h pattern then
          memoize_count_matches
            (String.sub pattern towel_len (pattern_len - towel_len))
            towels
          + has_pattern t
        else has_pattern t
  in
  match Hashtbl.find_opt memo2 pattern with
  | Some x -> x
  | None ->
      let x = if pattern_len = 0 then 1 else has_pattern towels in
      Hashtbl.add memo2 pattern x;
      x

let count_matches towels pattern = memoize_count_matches pattern towels

let part1 (towels, patterns) =
  List.filter (has_match towels) patterns |> List.length

let part2 (towels, patterns) =
  List.map (count_matches towels) patterns |> List.fold_left ( + ) 0

let _ =
  Aoc.main towels_of_file [ (string_of_int, part1); (string_of_int, part2) ]