module IntMap = Map.Make (Int)

(** [parse_robot s] returns a robot description [((x, y), (dx, dy))] parsed from
    the string [s]. *)
let parse_robot s =
  let re =
    Str.regexp {|p=\(-?[0-9]+\),\(-?[0-9]+\) v=\(-?[0-9]+\),\(-?[0-9]+\)|}
  in
  let _ = Str.search_forward re s 0 in
  ( ( int_of_string (Str.matched_group 1 s),
      int_of_string (Str.matched_group 2 s) ),
    ( int_of_string (Str.matched_group 3 s),
      int_of_string (Str.matched_group 4 s) ) )

(** [robots_of_file fname] returns a list of robots parsed from the file
    [fname]. *)
let robots_of_file fname = Aoc.strings_of_file fname |> List.map parse_robot

(** Grid width *)
let width = 101

(** Grid height *)
let height = 103

(** Number of seconds to run part 1 for*)
let secs1 = 100

(** Maximum number of seconds to run part 2 for *)
let secs2 = 1000000

(** [normalize_velocity robot] returns a robot where the velocity has been
    normalized to be non-negative in both directions. *)
let normalize_velocity (p, (dx, dy)) =
  (p, ((dx + width) mod width, (dy + height) mod height))

(** [calc_pos_after secs r] returns the [(x, y)] position of a robot after
    [secs] seconds. *)
let calc_pos_after secs ((x, y), (dx, dy)) =
  let x' = (x + (secs * dx)) mod width in
  let y' = (y + (secs * dy)) mod height in
  (x', y')

(** [in_a_quadrant pos] returns true if [pos] is in a quadrant. *)
let in_a_quadrant (x, y) = x <> width / 2 && y <> height / 2

(** [update_count n] Is used by [IntMap.update] to increment a count. *)
let update_count = function None -> Some 1 | Some x -> Some (x + 1)

(** [get_quadrant p] returns the quadrant ID that the position [p] is in. *)
let get_quadrant (x, y) =
  if x < width / 2 && y < height / 2 then 1
  else if x > width / 2 && y < height / 2 then 2
  else if x < width / 2 && y > height / 2 then 4
  else if x > width / 2 && y > height / 2 then 3
  else failwith "get_quadrant"

(** [quadrant_counts map p] updates the quadrant count map [map] with [p]. Keys
    to [map] are quadrant IDs, and the values are the number of robots in that
    quadrant. *)
let quadrant_counts map p =
  let idx = get_quadrant p in
  IntMap.update idx update_count map

(** [print_locs lst] prints the grid layout of the robots given in list. *)
let print_locs lst =
  let a = Array.make_matrix height width '.' in
  let rec impl = function
    | [] -> ()
    | (x, y) :: t ->
        if a.(y).(x) = '.' then a.(y).(x) <- '1'
        else a.(y).(x) <- char_of_int (1 + int_of_char a.(y).(x));
        impl t
  in
  impl lst;
  Array.iter
    (fun r ->
      Array.iter print_char r;
      print_newline ())
    a

(** [part1 robots] solves part1 for the list [robots]. *)
let part1 robots =
  let counts =
    robots
    |> List.map normalize_velocity
    |> List.map (calc_pos_after secs1)
    |> List.filter in_a_quadrant
    |> List.fold_left quadrant_counts IntMap.empty
  in
  IntMap.fold (fun _ v acc -> acc * v) counts 1

(** [find_tree max_n lst] tries to find the Christmas tree picture by iterating
    through the various steps robots move in. Returns the number of iterations
    before finding the tree. *)
let find_tree max_n lst =
  (* We assume that the picture will occur when every robot is in a unique
     location. *)
  let num_robots = List.length lst in
  let rec impl n =
    if n > max_n then failwith "None found"
    else
      let poses = List.map (calc_pos_after n) lst in
      (* If every tree is in a unique location then sort_uniq will not remove
         any elements from the list. *)
      if List.length (List.sort_uniq Aoc.IntPair.compare poses) = num_robots
      then (
        print_locs poses;
        n)
      else impl (n + 1)
  in
  impl 0

(** [part2 robots] solves part 2 for the list of robots. *)
let part2 robots =
  let robots = List.map normalize_velocity robots in
  find_tree secs2 robots

let _ =
  Aoc.main robots_of_file [ (string_of_int, part1); (string_of_int, part2) ]