let find_regions grid =
  let working = Array.make (Aoc.Grid.length grid) ~-1 in
  let add_pos region pos perimeter lst =
    if Aoc.Grid.pos_is_valid grid pos && Aoc.Grid.get_by_pos grid pos = region
    then (Aoc.Grid.idx_of_pos grid pos :: lst, perimeter - 1)
    else (lst, perimeter)
  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 + 4 in
        let region = Aoc.Grid.get_by_idx grid idx in
        let t, perimeter = add_pos region (x - 1, y) perimeter t in
        let t, perimeter = add_pos region (x + 1, y) perimeter t in
        let t, perimeter = add_pos region (x, y - 1) perimeter t in
        let t, perimeter = add_pos region (x, y + 1) perimeter 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 part1 grid =
  find_regions grid |> List.fold_left (fun acc (p, a) -> acc + (p * a)) 0

let _ = Aoc.main Aoc.Grid.of_file [ (string_of_int, part1) ]