module IntMap = Map.Make (Int)

let load_file fname =
  match In_channel.with_open_text fname In_channel.input_line with
  | Some x -> x
  | None -> failwith "load_file"

let rec apply_n n fn arg = if n <= 0 then arg else apply_n (n - 1) fn (fn arg)
let update_count n = function None -> Some n | Some x -> Some (x + n)

let map_of_ints =
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (IntMap.update h (update_count 1) acc) t
  in
  impl IntMap.empty

(** [calc_blink_rec acc lst] returns an updated map based off [acc] with the
    result of apply a blink step to the stones in [lst]. Entries in [lst] are
    pairs of [(stone id, count)]. [acc] and the resulting map have keys which
    are stone id, and values which are count. *)
let rec calc_blink_rec acc = function
  | [] -> acc
  | (0, n) :: t -> calc_blink_rec (IntMap.update 1 (update_count n) acc) t
  | (x, n) :: t when Aoc.digits10 x mod 2 = 0 ->
      let pow = Aoc.pow10 (Aoc.digits10 x / 2) in
      let acc = IntMap.update (x / pow) (update_count n) acc in
      let acc = IntMap.update (x mod pow) (update_count n) acc in
      calc_blink_rec acc t
  | (x, n) :: t ->
      calc_blink_rec (IntMap.update (x * 2024) (update_count n) acc) t

(** [calc_blink map] calculates how a collection of stones changes in a blink.
    [map] is a map with key of stone ID, and value number of times that stone
    appears. The result is a map with similar key, value pairs.

    This improves performance because we find that the transformation stones go
    through ends up producing repeated numbers. e.g.: 0 -> 1 -> 2024 -> 20 24 ->
    2 0 2 4 which has two 2s in it.

    We also note that despite the problem description saying stones stay in
    order, the result we are asked for (number of stones) does not require them
    to be in order. *)
let calc_blink map = IntMap.to_list map |> calc_blink_rec IntMap.empty

(** [part n str] returns the number of stones after [n] blinks, given an initial
    string [str] of space seperated stone IDs. *)
let part n str =
  let map = Aoc.ints_of_string str |> map_of_ints |> apply_n n calc_blink in
  IntMap.fold (fun _ v acc -> v + acc) map 0

let _ =
  Aoc.main load_file [ (string_of_int, part 25); (string_of_int, part 75) ]