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


class Day(AOCDay):
    inputs = [
        [
            (12, "input14_test"),
            (225648864, "input14"),
        ],
        [
            (7847, "input14"),
        ],
    ]

    def parse_input(self) -> list[tuple[tuple[int, int], tuple[int, int]]]:
        robots = []
        for line in self.getInput():
            p, v = line.split()
            px, py = map(int, p[2:].split(","))
            vx, vy = map(int, v[2:].split(","))
            robots.append(((px, py), (vx, vy)))

        return robots

    def move_robots(
        self, robots: list[tuple[tuple[int, int], tuple[int, int]]], steps: int = 100
    ) -> tuple[tuple[int, int], list[tuple[tuple[int, int], tuple[int, int]]]]:
        if self.is_test():
            dim = (11, 7)
        else:
            dim = (101, 103)

        return dim, [
            (((x[0][0] + steps * x[1][0]) % dim[0], (x[0][1] + steps * x[1][1]) % dim[1]), x[1]) for x in robots
        ]

    def part1(self) -> Any:
        q = [[0, 0, 0], [0, 0, 0], [0, 0, 0]]
        dim, robots = self.move_robots(self.parse_input())
        for robot in robots:
            q[compare(robot[0][1], dim[1] // 2) + 1][compare(robot[0][0], dim[0] // 2) + 1] += 1

        return q[0][0] * q[2][0] * q[0][2] * q[2][2]

    def part2(self) -> Any:
        robots = self.parse_input()

        count = 0
        while True:
            count += 1
            _, new_robots = self.move_robots(robots, steps=1)
            g = Grid()
            for robot in new_robots:
                g.set(Coordinate(robot[0][0], robot[0][1]))

            found = False
            for c in g.getActiveCells():
                if len(set(g.getRegion(c))) > 50:
                    found = True
                    break

            if found:
                break

            robots = new_robots

        return count


if __name__ == "__main__":
    day = Day(2024, 14)
    day.run(verbose=True)