from collections import defaultdict

import aoclib

DAY = 11
TEST_SOLUTION_PART1 = 37
TEST_SOLUTION_PART2 = 26
INT_INDEX = ['L', '#', '.']


def printLayout(layout, max_x, max_y):
    for y in range(0, max_y):
        for x in range(0, max_x):
            print(INT_INDEX[layout[(x, y)]], end="")

        print()


def prepareLayout(raw_input):
    int_layout = defaultdict(lambda: 2)
    for y, line in enumerate(raw_input):
        for x, char in enumerate(line):
            int_layout[(x, y)] = INT_INDEX.index(char)

    return int_layout


def simulate(layout, max_x, max_y, wide=False):
    num_seats_occupied = 0
    new_layout = defaultdict(lambda: 2)
    for x in range(0, max_x):
        for y in range(0, max_y):
            if layout[(x, y)] == 2:
                continue

            occupied_seats = 0
            for cx in [-1, 0, 1]:
                for cy in [-1, 0, 1]:
                    if cx == 0 and cy == 0:
                        continue
                    if not wide:
                        if layout[(x + cx, y + cy)] == 1:
                            occupied_seats += 1
                    else:
                        for d in range(1, max(max_x, max_y)):
                            if d * cy + y >= max_x or d * cx + x >= max_y:
                                break
                            if layout[(d * cx + x, d * cy + y)] == 2:
                                continue
                            if layout[(d * cx + x, d * cy + y)] < 2:
                                occupied_seats += layout[(d * cx + x, d * cy + y)]
                                break

            if layout[(x, y)] and occupied_seats >= (5 if wide else 4):
                new_layout[(x, y)] = 0
            elif not layout[(x, y)] and occupied_seats == 0:
                new_layout[(x, y)] = 1
            else:
                new_layout[(x, y)] = layout[(x, y)]

            num_seats_occupied += new_layout[(x, y)]

    return new_layout, num_seats_occupied


def part1(test_mode=False, wide=False):
    my_input = aoclib.getInputAs2DArray(day=DAY, test=test_mode)

    curr_layout = prepareLayout(my_input)
    cellcount_x = len(my_input[0])
    cellcount_y = len(my_input)
    curr_occupied = 0
    last_occupied = -1
    while curr_occupied != last_occupied:
        last_occupied = curr_occupied
        curr_layout, curr_occupied = simulate(curr_layout, cellcount_x, cellcount_y, wide)

    return curr_occupied


def part2(test_mode=False):
    return part1(test_mode, True)