open Printf;;
open Buffer;;
open Str;;

let read_whole_file f =
  let s = really_input_string f (in_channel_length f) in
  close_in f;
  s

let rec find_all_numbers pos str  =
  try
    let re = Str.regexp "-?[0-9]+" in
    let new_pos = Str.search_forward re str pos in
    let n_str = Str.matched_group 0 str in
    let n = int_of_string n_str in
    n :: find_all_numbers (new_pos + String.length n_str) str 
  with Not_found -> []

let rec to_quadruples lst =
  match lst with
  | n1 :: n2 :: n3 :: n4 :: tail ->
    (n1, n2, n3, n4) :: to_quadruples tail
  | _ -> []

let true_modulo a b =
  ((a mod b) + b) mod b

let find_state_after_steps (upper_x, upper_y) steps (x, y, vx, vy) =
  (true_modulo (x + vx * steps) upper_x, true_modulo (y + vy * steps) upper_y, vx, vy)

let transform_all_states (upper_x, upper_y) steps states_list =
  List.map (find_state_after_steps (upper_x, upper_y) steps) states_list

let count_in_area (lower_x, lower_y) (upper_x, upper_y) states_list =
  states_list
    |> List.filter (fun (x, y, _, _) -> lower_x <= x && x < upper_x && lower_y <= y && y < upper_y)
    |> List.length

let get_safety (upper_x, upper_y) states_list =
  count_in_area (0, 0) (upper_x / 2, upper_y / 2) states_list
    * count_in_area (upper_x / 2 + 1, 0) (upper_x, upper_y / 2) states_list
    * count_in_area (0, upper_y / 2 + 1) (upper_x / 2, upper_y) states_list
    * count_in_area (upper_x / 2 + 1, upper_y / 2 + 1) (upper_x, upper_y) states_list

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 rec count_cluster occupied_set (x, y) =
  match IntPairsSet.find_opt (x, y) occupied_set with
  | Some _ ->
    [(x - 1, y); (x, y - 1); (x + 1, y); (x, y + 1)]
      |> List.fold_left (
          fun (acc_count, acc_set) pos -> 
            let inner_count, new_set = count_cluster acc_set pos in
            (acc_count + inner_count, new_set)
        )
        (1, IntPairsSet.remove (x, y) occupied_set)
  | None -> (0, occupied_set)

let rec find_all_clusters_inner occupied_set =
  match IntPairsSet.min_elt_opt occupied_set with
  | Some root ->
    let count, new_set = count_cluster occupied_set root in
    count :: find_all_clusters_inner new_set
  | None -> []

let find_all_clusters (upper_x, upper_y) state_list =
  let occupied = List.fold_left (fun set (x, y, _, _) -> IntPairsSet.add (x, y) set) IntPairsSet.empty state_list in
  find_all_clusters_inner occupied
  |> List.sort (fun a b -> Stdlib.compare b a)

let print_field (upper_x, upper_y) state_list =
  let arr = Array.make_matrix upper_y upper_x '.' in
  List.iter (fun (x, y, _, _) -> arr.(y).(x) <- '#') state_list;
  arr
    |> Array.iter (fun row -> row
      |> Array.iter (fun x -> print_char x);
      print_endline "")

let is_probably_a_tree_with_border cluster_sizes =
  match cluster_sizes with
  | _ :: n2 :: tail -> begin
    match tail with
    | n3 :: _ -> n3 * 20 < n2
    | _ -> true
  end
  | _ -> false

let rec search_for_a_tree (upper_x, upper_y) state_list i =
  let cluster_sizes = find_all_clusters (upper_x, upper_y) state_list in
  if is_probably_a_tree_with_border cluster_sizes then
    i
  else if i > upper_x * upper_y then
    -1
  else
    search_for_a_tree (upper_x, upper_y) (transform_all_states (upper_x, upper_y) 1 state_list) (i + 1)

let () =
  let f = open_in_bin "input.txt" in
  let upper_x = 101 in
  let upper_y = 103 in
  let robots = f
    |> read_whole_file
    |> find_all_numbers 0
    |> to_quadruples in
  let final_states = transform_all_states (upper_x, upper_y) 100 robots in
  let result = get_safety (upper_x, upper_y) final_states in
  let result2 = search_for_a_tree (upper_x, upper_y) robots 0 in
  print_field (upper_x, upper_y) (transform_all_states (upper_x, upper_y) result2 robots);
  printf "%d\n%d\n" result result2