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_get_pos map pos = map_get_by_idx map (map_idx_of_pos map pos)*)

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

(*let map_is_valid_idx map idx = idx >= 0 && idx < map_length map*)

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 get_antinodes map station_idxs =
  let rec impl2 acc (px, py) = function
    | [] -> acc
    | h :: t ->
        let hx, hy = map_pos_of_idx map h in
        let dx = hx - px in
        let dy = hy - py in
        impl2 ((px - dx, py - dy) :: (hx + dx, hy + dy) :: acc) (px, py) t
  in
  let rec impl acc = function
    | [] -> acc
    | h :: t -> impl (impl2 acc (map_pos_of_idx map h) t) t
  in
  let lst = impl [] station_idxs |> List.filter (map_is_valid_pos map) in
  let lst2 = lst |> List.map (map_idx_of_pos map) in
  lst2

let process_stations map stations = Array.map (get_antinodes map) stations

let string_replace str pos c =
  assert (pos < String.length str);
  String.sub str 0 pos ^ String.make 1 c
  ^ String.sub str (pos + 1) (String.length str - pos - 1)

let map_print map =
  let rec impl off = if off < map_length map then impl (off + map.width) in
  impl 0

let rec add_stations map = function
  | [] -> map
  | h :: t -> add_stations { map with map = string_replace map.map h '#' } t

let part1 map =
  let lst =
    get_station_indices map |> process_stations map
    |> Array.fold_left List.append []
    |> List.sort_uniq Stdlib.compare
  in
  let map' = add_stations map lst in
  map_print map';
  List.length lst

let _ = Aoc.main map_of_file [ (string_of_int, part1) ]