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)