from collections import defaultdict

from tools.aoc import AOCDay
from tools.coordinate import Coordinate
from tools.grid import Grid
from typing import Any


CHECK_ORDER = [
    (Coordinate(0, -1), Coordinate(-1, -1), Coordinate(1, -1)),
    (Coordinate(0, 1), Coordinate(1, 1), Coordinate(-1, 1)),
    (Coordinate(-1, 0), Coordinate(-1, 1), Coordinate(-1, -1)),
    (Coordinate(1, 0), Coordinate(1, 1), Coordinate(1, -1))
]


def move_elfs(map: Grid, round: int) -> bool:
    order_index = round % 4

    move_to = {}
    proposed_positions = defaultdict(int)
    for elf in map.getActiveCells():
        if map.getNeighbourSum(elf, includeDiagonal=True) == 0:
            continue

        for i in range(4):
            check_dir = CHECK_ORDER[(order_index + i) % 4]
            free = True
            for check_pos in check_dir:
                if map.get(elf + check_pos):
                    free = False
                    break

            if free:
                move_to[elf] = elf + check_dir[0]
                proposed_positions[elf + check_dir[0]] += 1
                break

    if not move_to:
        return False

    for elf, target in move_to.items():
        if proposed_positions[target] == 1:
            map.toggle(elf)
            map.toggle(target)

    return True


class Day(AOCDay):
    inputs = [
        [
            (25, "input23_test_small"),
            (110, "input23_test"),
            (3766, "input23"),
        ],
        [
            (20, "input23_test"),
            (954, "input23"),
        ]
    ]

    def part1(self) -> Any:
        map = Grid.from_str("/".join(self.getInput()), translate={'#': True, '.': False})

        for i in range(10):
            if not move_elfs(map, i):
                break

        map.recalcBoundaries()
        return (map.maxX - map.minX + 1) * (map.maxY - map.minY + 1) - map.getOnCount()

    def part2(self) -> Any:
        map = Grid.from_str("/".join(self.getInput()), translate={'#': True, '.': False})
        round = 0

        while move_elfs(map, round):
            round += 1

        return round + 1


if __name__ == '__main__':
    day = Day(2022, 23)
    day.run(verbose=True)