from collections import defaultdict

import aoclib

DAY = 11
TEST_SOLUTION_PART1 = 37
TEST_SOLUTION_PART2 = 26
INT_INDEX = ['L', '#', '.']
MAP_INDEX = [False, True, None]


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: None)
    for y, line in enumerate(raw_input):
        for x, char in enumerate(line):
            int_layout[(x, y)] = MAP_INDEX[INT_INDEX.index(char)]

    return int_layout


def simulate(layout, max_x, max_y, wide=False):
    max_seats_occupied = (5 if wide else 4)
    new_layout = defaultdict(lambda: None)
    for x in range(0, max_x):
        for y in range(0, max_y):
            if layout[(x, y)] is None:
                continue

            occupied_seats = 0
            for cx in [x-1, x, x+1]:
                for cy in [y-1, y, y+1]:
                    if cx == x and cy == y:
                        continue

                    check_value = layout[(cx, cy)]

                    if check_value is not None:
                        occupied_seats += check_value
                    elif wide:
                        test_x = cx + (cx - x)
                        test_y = cy + (cy - y)
                        while 0 <= test_x <= max_x and 0 <= test_y <= max_y:
                            if layout[(test_x, test_y)] is None:
                                test_x += (cx - x)
                                test_y += (cy - y)
                            else:
                                occupied_seats += layout[(test_x, test_y)]
                                break

            if occupied_seats >= max_seats_occupied:
                new_layout[(x, y)] = False
            elif occupied_seats == 0:
                new_layout[(x, y)] = True
            else:
                new_layout[(x, y)] = layout[(x, y)]

    return new_layout


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
    generation_counter = 0

    # Yes, this is relying on the assumption that there are no
    # two different patterns with the same number of occupants.
    while curr_occupied != last_occupied:
        last_occupied = curr_occupied
        generation_counter += 1
        curr_layout = simulate(curr_layout, cellcount_x, cellcount_y, wide)
        curr_occupied = sum(filter(None, curr_layout.values()))

    print("Generations simulated: %d" % generation_counter)
    return curr_occupied


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