#!/usr/bin/env python

https://adventofcode.com/2025/day/4

f = open("day04input.txt", "r")
#f = open("testinput.txt", "r")
d = [l.strip() for l in f.readlines()]


def day04(data: list[str], remove_rolls: bool = False) -> int:
  row_no = len(data)
  col_no = len(data[0])
  data_list = list("".join(data))
  assert row_no * col_no == len(data_list)

  offsets = {-col_no - 1, -col_no, -col_no + 1, -1, 1, col_no - 1, col_no, col_no + 1}
  offset_l = {-col_no - 1, -1, col_no - 1}
  offset_r = {-col_no + 1, 1, col_no + 1}
  offset_t = {-col_no - 1, -col_no, -col_no + 1}
  offset_b = {col_no - 1, col_no, col_no + 1}

  # little debug function that prints the current floor plan
  def field_to_str(data_list: list) -> str:
    field = ""
    for j in range(col_no):
      field += "".join(data_list[j*col_no:(j+1)*col_no-1]) + "\n"
    return field

#  print(f"initial state")
#  print(field_to_str(data_list))

  def count_rolls(data_list: list) -> int:
    iteration_rolls = 0

    # loop over data 
    for i in range(len(data_list)):
      # skip if there's no roll on this position
      if data_list[i] == ".":
        continue
      
      local_offsets = offsets.copy()

      # check left edge 
      if i % col_no == 0:
        local_offsets -= offset_l
      # check right edge 
      if i % col_no == col_no - 1:
        local_offsets -= offset_r
      # check top edge 
      if i - col_no < 0:
        local_offsets -= offset_t
      # check bottom edge 
      if i + col_no >= len(data_list):
        local_offsets -= offset_b
      assert len(local_offsets) in (3, 5, 8)

      # loop over adjacent positions 
      adjacent_rolls = 0
      for o in local_offsets:
        # adjacent position is empty
        if (c := data_list[i + o]) == ".":
          continue
        # adjacent position is a roll 
        elif c in  "x@":
          adjacent_rolls += 1
        # adjacent position is unknown character
        else:
          raise Exception("arrgh!")

      # mark roll as removable
      if adjacent_rolls < 4:
        iteration_rolls += 1
        data_list[i] = "x"

    return iteration_rolls


  # solve part 1
  if not remove_rolls:
    return count_rolls(data_list)

  # solve part 2
  total_rolls = 0
  while True:
    iteration_rolls = count_rolls(data_list)
    if iteration_rolls > 0:
      # increase number of removed rolls
      total_rolls += iteration_rolls
      # remove marked rolls
      data_list = list(map(lambda c: "." if c == "x" else c, data_list))
    else:
      return total_rolls
  

print(day04(d))
print(day04(d, True))