let get_by_pos_opt grid pos =
  if Aoc.Grid.pos_is_valid grid pos then Some (Aoc.Grid.get_by_pos grid pos)
  else None

let perimeter grid (x, y) =
  let region = Aoc.Grid.get_by_pos grid (x, y) in
  0
  + (if get_by_pos_opt grid (x - 1, y) = Some region then 1 else 0)
  + (if get_by_pos_opt grid (x + 1, y) = Some region then 1 else 0)
  + (if get_by_pos_opt grid (x, y + 1) = Some region then 1 else 0)
  + if get_by_pos_opt grid (x, y - 1) = Some region then 1 else 0

let is_corner grid (x, y) (dx, dy) =
  let region = Aoc.Grid.get_by_pos grid (x, y) in
  if
    get_by_pos_opt grid (x + dx, y) <> Some region
    && get_by_pos_opt grid (x, y + dy) <> Some region
  then true
  else if
    get_by_pos_opt grid (x + dx, y) = Some region
    && get_by_pos_opt grid (x, y + dy) = Some region
    && get_by_pos_opt grid (x + dx, y + dy) <> Some region
  then true
  else false

let corners grid pos =
  (if is_corner grid pos (-1, -1) then 1 else 0)
  + (if is_corner grid pos (-1, 1) then 1 else 0)
  + (if is_corner grid pos (1, 1) then 1 else 0)
  + if is_corner grid pos (1, -1) then 1 else 0

let find_regions calc grid =
  let working = Array.make (Aoc.Grid.length grid) ~-1 in
  let add_pos region pos lst =
    if 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 id = function
    | [] -> (perimeter, area)
    | idx :: t when working.(idx) <> -1 -> scan_pos perimeter area id t
    | idx :: t ->
        working.(idx) <- id;
        let x, y = Aoc.Grid.pos_of_idx grid idx in
        let area = succ area 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 id t
  in
  let rec impl acc idx id =
    if idx >= Aoc.Grid.length grid then List.rev acc
    else
      let perimeter, area = scan_pos 0 0 id [ idx ] in
      if area = 0 then impl acc (idx + 1) id
      else impl ((perimeter, area) :: acc) (idx + 1) (id + 1)
  in
  impl [] 0 0

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) ]