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 file_to_2d_array in_file =
  let lines = list_of_lines in_file in
  let size_x = String.length @@ List.hd lines in
  let size_y = List.length lines in
  let arr = Array.make_matrix size_x size_y '?' in
  List.iteri (fun j line -> String.iteri (fun i c -> arr.(j).(i) <- c) line) lines;
  arr

module CharMap = Map.Make(Char)
module IntPairs =
struct
  type t = int * int
  let compare (x0,y0) (x1,y1) =
    match Stdlib.compare x0 x1 with
        0 -> Stdlib.compare y0 y1
      | c -> c
end
module IntPairsSet = Set.Make(IntPairs)

let all_objects_in_row arr =
  arr |> Array.to_seq
      |> Seq.mapi (fun i x -> (i, x))
      |> Seq.filter (fun (i, x) -> x <> '.')

let update_fun elem prev =
  match prev with 
  | Some lst -> Some (elem :: lst)
  | None -> Some [elem]

let group_coords_by_elem arr =
  arr |> Array.to_seq
      |> Seq.mapi (fun j row -> Seq.map (fun (i, c) -> ((i, j), c)) (all_objects_in_row row))
      |> Seq.concat
      |> Seq.fold_left (fun acc (pos, c) -> CharMap.update c (update_fun pos) acc) CharMap.empty

let in_bounds (upper_x, upper_y) (x, y) =
  0 <= x && x < upper_x && 0 <= y && y < upper_y

let get_antinodes_for_pair (x1, y1) (x2, y2) =
  let diff_x = x2 - x1 in
  let diff_y = y2 - y1 in
  [(x1 - diff_x, y1 - diff_y); (x2 + diff_x, y2 + diff_y)]

let rec get_antinodes_pairwise antinode_generator antennae_list =
  match antennae_list with
  | head :: tail -> 
    (List.fold_left (fun acc x -> List.append acc (antinode_generator head x)) [] tail)
    |> List.append (get_antinodes_pairwise antinode_generator tail)
  | _ -> []

let rec extend_series upper_bounds diff start =
  if not @@ in_bounds upper_bounds start then
    []
  else
    start :: extend_series upper_bounds diff (fst start + fst diff, snd start + snd diff)

let get_harmonic_antinodes_for_pair upper_bounds (x1, y1) (x2, y2) =
  let diff_x = x2 - x1 in
  let diff_y = y2 - y1 in
  List.append 
    (extend_series upper_bounds (-diff_x, -diff_y) (x1, y1))
    (extend_series upper_bounds (diff_x, diff_y) (x2, y2))
  

let () =
  let f = open_in "input.txt" in
  let arr = file_to_2d_array f in
  let antennae_grouped = group_coords_by_elem arr in
  let upper_bounds = (Array.length arr.(0), Array.length arr) in
  let antinode_set = antennae_grouped
    |> CharMap.to_seq
    |> List.of_seq
    |> List.map snd
    |> List.map (get_antinodes_pairwise (fun a b -> get_antinodes_for_pair a b
                                        |> List.filter (in_bounds upper_bounds)) )
    |> List.concat
    |> IntPairsSet.of_list in
  let harmonic_antinode_set = antennae_grouped
    |> CharMap.to_seq
    |> List.of_seq
    |> List.map snd
    |> List.map (get_antinodes_pairwise (fun a b -> get_harmonic_antinodes_for_pair upper_bounds a b))
    |> List.concat
    |> IntPairsSet.of_list in
  let result = IntPairsSet.cardinal antinode_set in
  let result2 = IntPairsSet.cardinal harmonic_antinode_set in
  printf "%d\n%d\n" result result2