FLOOR = '.'
EMPTY = 'L'
OCCUPIED = '#'

def n_occupied(seat: str) -> int:
    ''' Returns 1 if occupied and 0 if not. '''
    return 1 if seat == OCCUPIED else 0


def compute_next_gen(grid: str, stride: int) -> str:
    next_gen = []
    for i, seat in enumerate(grid):
        if seat == FLOOR:
            next_gen.append(FLOOR)
            continue
        # Compute occupied adjacents; add left, right, above and below.
        occupied_neighbors = n_occupied(grid[i - 1]) \
                           + n_occupied(grid[i + 1]) \
                           + n_occupied(grid[i - stride]) \
                           + n_occupied(grid[i + stride])
        if n_occupied(seat) == 0:
            next_gen.append(OCCUPIED if occupied_neighbors == 0 else EMPTY)
        else:
            next_gen.append(OCCUPIED if occupied_neighbors < 4 else EMPTY)
    return ''.join(next_gen)


def compute_steady_state(grid: str, width: int) -> str:
    # Make sure the grid is well formed.
    assert len(grid) % width == 0

    stride = width + 2
    # Surround the grid with empties so we don't have to do special casing for
    # edges in copmute_next_gen. Instead we could use a real position and
    working_grid = []
    working_grid.extend(FLOOR * stride) # Top row
    for row_span_start in range(0, len(grid), width):
        row_span_end = row_span_start + width
        working_grid.append(FLOOR)
        working_grid.append(grid[row_span_start:row_span_end])
        working_grid.append(FLOOR)
    # Bottom row
    working_grid.extend(FLOOR * stride) # Bottom row

    # Make it a string
    working_grid = ''.join(working_grid)
    next_state = None

    while True:
        next_state = compute_next_gen(working_grid, stride)
        if working_grid == next_state:
            break
        working_grid = next_state

    steady_state = []
    # Back to the og grid
    for row_span_start in range(stride + 1, len(working_grid) - stride, stride):
        steady_state.extend(working_grid[row_span_start:row_span_start + width])

    steady_state = ''.join(steady_state)
    return steady_state


def main():
    grid = \
          'L.LL.LL.LL' \
        + 'LLLLLLL.LL' \
        + 'L.L.L..L..' \
        + 'LLLL.LL.LL' \
        + 'L.LL.LL.LL' \
        + 'L.LLLLL.LL' \
        + '..L.L.....' \
        + 'LLLLLLLLLL' \
        + 'L.LLLLLL.L' \
        + 'L.LLLLL.LL'
    print(grid)
    width = 10

    steady_state = compute_steady_state(grid, width)
    print("Steady state:")
    for span_start in range(0, len(steady_state), width):
        span_end = span_start + width
        print(steady_state[span_start:span_end])

if __name__ == '__main__':
    main()