module CharMap = Map.Make (Char)

(** [get_station_indices map] returns a list of pairs mapping station ID to the
    indices in [map] where there is a station with that ID. *)
let get_station_indices map =
  let rec impl acc idx =
    if idx >= Aoc.Grid.length map then acc
    else if Aoc.Grid.get_by_idx map idx = '.' then impl acc (idx + 1)
    else
      let station = Aoc.Grid.get_by_idx map idx in
      let update_fn lst =
        match lst with None -> Some [ idx ] | Some t -> Some (idx :: t)
      in
      impl (CharMap.update station update_fn acc) (idx + 1)
  in
  impl CharMap.empty 0 |> CharMap.to_list

(** Generate antinodes for part 1. *)
let get_antinodes1 acc _ (px, py) (px', py') =
  let dx = px' - px in
  let dy = py' - py in
  (px - dx, py - dy) :: (px' + dx, py' + dy) :: acc

(** [add_antinodes lst map pos vel] adds antinodes at [pos + n * vel] to [lst]
    for all non-negative [n] that are valid positions in [map]. *)
let rec add_antinodes lst map (x, y) (dx, dy) =
  if Aoc.Grid.pos_is_valid map (x, y) then
    add_antinodes ((x, y) :: lst) map (x + dx, y + dy) (dx, dy)
  else lst

(** Generate antinodes for part 2. *)
let get_antinodes2 acc map (px, py) (px', py') =
  let dx = px' - px in
  let dy = py' - py in
  let acc' = add_antinodes acc map (px, py) (dx, dy) in
  let acc'' = add_antinodes acc' map (px, py) (-dx, -dy) in
  acc''

(** [process_stations map fn stations] generates a list of all antinodes for the
    stations in [stations] on the map [map]. [fn acc map p p'] is called to
    generate the antinode list for each pair of stations [p] and [p']. It should
    add the positions of antinodes to the list [acc]. *)
let process_stations map fn stations =
  let rec impl2 acc p t =
    match t with
    | [] -> acc
    | h :: t -> impl2 (fn acc map p (Aoc.Grid.pos_of_idx map h)) p t
  in
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (impl2 acc (Aoc.Grid.pos_of_idx map h) t) t
  in
  List.map (impl []) stations

(** [part antifn map station_indices] process all the stations in
    station_indices calling [antifn acc map p p'] on all stations. Here [acc] is
    a list of antinodes which [antifn] should update and return, [p] and [p']
    are positions of stations to generate antinodes for. *)
let part antifn map =
  get_station_indices map
  |> List.map snd (* we do not care about the station IDs *)
  |> process_stations map antifn
  |> List.concat
  |> List.filter (Aoc.Grid.pos_is_valid map)
  |> List.sort_uniq Stdlib.compare
  |> List.length

let _ =
  Aoc.main Aoc.Grid.of_file
    [
      (string_of_int, part get_antinodes1); (string_of_int, part get_antinodes2);
    ]