(** [perimeter grid pos] returns the number of sides of [pos] that contribute to
    the perimeter of the region that [pos] is part of in [grid]. *)
let perimeter grid (x, y) =
  let region = Aoc.Grid.get_by_pos grid (x, y) in
  let check (dx, dy) =
    Aoc.Grid.get_by_pos_opt grid (x + dx, y + dy) = Some region
  in
  [ (-1, 0); (1, 0); (0, 1); (0, -1) ]
  |> List.map check |> List.map Bool.to_int |> List.fold_left ( + ) 0

(** [is_corner grid pos dir] returns true if there is a region corner on the
    [dir] side of [pos] in grid. *)
let is_corner grid (x, y) (dx, dy) =
  let region = Aoc.Grid.get_by_pos grid (x, y) in
  if
    Aoc.Grid.get_by_pos_opt grid (x + dx, y) <> Some region
    && Aoc.Grid.get_by_pos_opt grid (x, y + dy) <> Some region
  then true
  else if
    Aoc.Grid.get_by_pos_opt grid (x + dx, y) = Some region
    && Aoc.Grid.get_by_pos_opt grid (x, y + dy) = Some region
    && Aoc.Grid.get_by_pos_opt grid (x + dx, y + dy) <> Some region
  then true
  else false

(** [corners grid pos] returns the count of the number of corners that [pos] is
    on for its region in [grid]. *)
let corners grid pos =
  [ (-1, -1); (-1, 1); (1, 1); (1, -1) ]
  |> List.map (is_corner grid pos)
  |> List.map Bool.to_int |> List.fold_left ( + ) 0

(** [find_regions calc grid] finds all the regions in [grid] and returns a list
    containing an element for each region. The element is a pair [(p, a)] where
    [p] is the sum of calling [calc grid pos] for every element in the region
    and [a] is the area of the region. *)
let find_regions calc grid =
  let visited = Array.make (Aoc.Grid.length grid) false in
  let add_pos region pos lst =
    if Aoc.Grid.get_by_pos_opt grid pos = Some region then
      Aoc.Grid.idx_of_pos grid pos :: lst
    else lst
  in
  let rec scan_pos perimeter area = function
    | [] -> (perimeter, area)
    | idx :: t when visited.(idx) -> scan_pos perimeter area t
    | idx :: t (* when not visited.(idx) *) ->
        visited.(idx) <- true;
        let x, y = Aoc.Grid.pos_of_idx grid idx in
        let area = area + 1 in
        let perimeter = perimeter + calc grid (x, y) in
        let region = Aoc.Grid.get_by_idx grid idx in
        let t = add_pos region (x - 1, y) t in
        let t = add_pos region (x + 1, y) t in
        let t = add_pos region (x, y - 1) t in
        let t = add_pos region (x, y + 1) t in
        scan_pos perimeter area t
  in
  let rec impl acc idx =
    if idx >= Aoc.Grid.length grid then List.rev acc
    else
      let perimeter, area = scan_pos 0 0 [ idx ] in
      if area = 0 then impl acc (idx + 1)
      else impl ((perimeter, area) :: acc) (idx + 1)
  in
  impl [] 0

(** [part calc grid] returns the result of Part N over [grid] where [calc] is
    the perimeter/edge counting function. *)
let part calc grid =
  find_regions calc grid |> List.fold_left (fun acc (p, a) -> acc + (p * a)) 0

let _ =
  Aoc.main Aoc.Grid.of_file
    [ (string_of_int, part perimeter); (string_of_int, part corners) ]