type map = { map : string; width : int; height : int }

let map_of_file fname =
  let strs = Aoc.strings_of_file fname in
  let width = String.length (List.hd strs) in
  let map = List.fold_left ( ^ ) "" strs in
  let height = String.length map / width in
  { map; width; height }

let map_length map = String.length map.map
let map_get_by_idx map idx = map.map.[idx]
let map_pos_of_idx map idx = (idx mod map.width, idx / map.width)
let map_idx_of_pos map (x, y) = x + (y * map.width)

let map_is_valid_pos map (x, y) =
  x >= 0 && x < map.width && y >= 0 && y < map.height

let get_station_indices map =
  let arr = Array.make 256 [] in
  let rec impl idx =
    if idx >= map_length map then arr
    else if map_get_by_idx map idx = '.' then impl (idx + 1)
    else
      let station = int_of_char (map_get_by_idx map idx) in
      arr.(station) <- idx :: arr.(station);
      impl (idx + 1)
  in
  impl 0

let rec get_antinodes1 acc map p =
  let px, py = map_pos_of_idx map p in
  function
  | [] -> acc
  | h :: t ->
      let hx, hy = map_pos_of_idx map h in
      let dx = hx - px in
      let dy = hy - py in
      get_antinodes1 ((px - dx, py - dy) :: (hx + dx, hy + dy) :: acc) map p t

let rec add_antinodes lst map (x, y) (dx, dy) =
  if map_is_valid_pos map (x, y) then
    add_antinodes ((x, y) :: lst) map (x + dx, y + dy) (dx, dy)
  else lst

let rec get_antinodes2 acc map p =
  let px, py = map_pos_of_idx map p in
  function
  | [] -> acc
  | h :: t ->
      let hx, hy = map_pos_of_idx map h in
      let dx = hx - px in
      let dy = hy - py in
      let acc' = add_antinodes acc map (px, py) (dx, dy) in
      let acc'' = add_antinodes acc' map (px, py) (-dx, -dy) in
      get_antinodes2 acc'' map p t

let process_stations map fn stations =
  let rec impl acc = function [] -> acc | h :: t -> impl (fn acc map h t) t in
  Array.map
    (fun x ->
      impl [] x
      |> List.filter (map_is_valid_pos map)
      |> List.map (map_idx_of_pos map))
    stations

let part antifn map =
  get_station_indices map
  |> process_stations map antifn
  |> Array.fold_left List.append []
  |> List.sort_uniq Stdlib.compare
  |> List.length

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