open Printf;;
open Map;;

let rec list_of_lines in_file =
  try
    let line = input_line in_file in
    line :: list_of_lines(in_file)
  with End_of_file ->
    close_in in_file;
    []

let rec parse_int_list sep str =
  let list1 = (List.filter (fun a -> String.length a != 0) (String.split_on_char sep str)) in
  List.map int_of_string list1

let stone_tick value = 
  if value = 0 then
    [1]
  else
    let digits_str = string_of_int value in
    let num_digits = String.length digits_str in
    if num_digits mod 2 = 0 then
      let half_num_digits = num_digits / 2 in
      let n1 = int_of_string (String.sub digits_str 0 half_num_digits) in
      let n2 = int_of_string (String.sub digits_str half_num_digits half_num_digits) in
      [n1; n2]
    else
      [value * 2024]

let transform_stones_naive stones = 
  stones
    |> List.map stone_tick
    |> List.concat

module IntMap = Map.Make(Int)

let add_counter_in_map n prev_opt =
  match prev_opt with 
  | Some prev_count -> Some (prev_count + n)
  | None -> Some n

let update_counter_in_map_accumulator n acc_map key =
  IntMap.update key (add_counter_in_map n) acc_map

let transform_stone_with_count stone_value count acc_map =
  let stone_result = stone_tick stone_value in
  List.fold_left (update_counter_in_map_accumulator count) acc_map stone_result

let transform_stones stone_counts_map =
  IntMap.fold transform_stone_with_count stone_counts_map IntMap.empty

let total_counters_in_map map =
  IntMap.fold (fun _ counter acc -> acc + counter) map 0

let rec apply_n_times n f arg =
  if n <= 0 then
    arg
  else
    apply_n_times (n-1) f (f arg)

let () =
  let f = open_in "input.txt" in
  let stones = f
    |> list_of_lines
    |> List.hd
    |> parse_int_list ' ' in
let result = stones
  |> apply_n_times 25 transform_stones_naive
  |> List.length in
printf "%d\n" result;
let result2 = stones
  |> List.fold_left (update_counter_in_map_accumulator 1) IntMap.empty 
  |> apply_n_times 75 transform_stones
  |> total_counters_in_map in
printf "%d\n" result2