(executable
  (name solution))

(lang dune 2.9)

NNCB
CH -> B
HH -> N
CB -> H
NH -> C
HB -> C
HC -> B
HN -> C
NN -> C
BH -> H
NC -> B
NB -> B
BN -> B
BB -> N
BC -> B
CC -> N
CN -> C

let read_lines path =
  let lines = ref [] in
  let chan = open_in path in
  try
    while true; do
      lines := input_line chan :: !lines
    done; !lines
  with End_of_file ->
    close_in chan;
    List.rev !lines
let h_mod tbl k f default =
  let new_value =
    if Hashtbl.mem tbl k
    then Hashtbl.find tbl k |> f
    else default
  in Hashtbl.replace tbl k new_value
let parse input =
  let template = List.hd input in
  let mappings = input |> List.tl |> List.tl in
  let counts = Hashtbl.create (List.length mappings) in
  let table = Hashtbl.create (List.length mappings) in
  let h_template = Hashtbl.create (String.length template) in
  let add_mapping line =
    Hashtbl.add table (line.[0], line.[1]) line.[6]
  in
  for i = 0 to (String.length template) - 2 do
    h_mod h_template (template.[i], template.[i+1]) ((+) 1) 1;
  done;
  String.iter (fun c -> h_mod counts c ((+) 1) 1) template;
  List.iter add_mapping mappings;
  (h_template, counts, table)
let step pairs counts mappings =
  let tmp_pairs = Hashtbl.copy pairs in
  let expand_pair pair count =
    let (a, b) = pair in
    let next = Hashtbl.find mappings pair in
    h_mod counts next ((+) count) count;
    h_mod pairs (a, next) ((+) count) count;
    h_mod pairs (next, b) ((+) count) count;
    h_mod pairs pair (fun x -> x - count) 0;
  in Hashtbl.iter expand_pair tmp_pairs
let rec times n f =
  if n > 0
  then (f (); times (n - 1) f)
let list_agg f l =
  List.fold_left f (List.hd l) l
let solution steps template counts mappings =
  let counts = Hashtbl.copy counts in
  let template = Hashtbl.copy template in
  times steps (fun () -> step template counts mappings);
  let value_counts = counts |> Hashtbl.to_seq_values |> List.of_seq in
  let highest = list_agg max value_counts in
  let lowest = list_agg min value_counts in
  highest - lowest
let part1 = solution 10
let part2 = solution 40
let input = read_lines Sys.argv.(1)
let (h_template, counts, mappings) = parse input
let () = Printf.printf "Part 1: %d\n" (part1 h_template counts mappings)
let () = Printf.printf "Part 2: %d\n" (part2 h_template counts mappings)