import re
from collections import Counter
from itertools import count
from math import prod
from textwrap import dedent
from typing import NamedTuple, override
from unittest import TestCase

from more_itertools import filter_map, first_true, iterate

from puzzles._solver import Solver


class Robot(NamedTuple):
    x: int
    y: int
    dx: int
    dy: int


robot_spec_pattern = re.compile(r"p=(-?\d+),(-?\d+) v=(-?\d+),(-?\d+)")


class DayFourteenSolver(Solver):
    robots: list[Robot]
    width: int
    height: int

    @override
    def __init__(self, puzzle_input: str, width=101, height=103):
        self.robots = [
            Robot(*map(int, match.groups()))
            for match in robot_spec_pattern.finditer(puzzle_input)
        ]
        self.width = width
        self.height = height

    @override
    def solve_p1(self) -> int:
        quadrant_counts = Counter(
            filter_map(
                lambda r: self.get_quadrant(self.move_robot(r, 100)),
                self.robots,
            )
        )
        return prod(quadrant_counts.values())

    @override
    def solve_p2(self) -> int:
        state = iterate(
            lambda robots: [self.move_robot(robot, 1) for robot in robots],
            self.robots,
        )
        return first_true(
            count(),
            pred=lambda i: self.total_overlap(next(state)) == 0,
        )

    def move_robot(self, robot: Robot, seconds: int) -> Robot:
        return Robot(
            (robot.x + robot.dx * seconds) % self.width,
            (robot.y + robot.dy * seconds) % self.height,
            robot.dx,
            robot.dy,
        )

    def get_quadrant(self, robot: Robot) -> int | None:
        if robot.x > self.width // 2 and robot.y < self.height // 2:
            return 1
        if robot.x < self.width // 2 and robot.y < self.height // 2:
            return 2
        if robot.x < self.width // 2 and robot.y > self.height // 2:
            return 3
        if robot.x > self.width // 2 and robot.y > self.height // 2:
            return 4
        return None

    def total_overlap(self, robots: list[Robot]) -> int:
        locations = Counter((robot.x, robot.y) for robot in robots)
        return sum(count for count in locations.values() if count > 1)


class TestDayFourteenSolver(TestCase):
    def test(self):
        solver = DayFourteenSolver(
            dedent(
                """
                p=0,4 v=3,-3
                p=6,3 v=-1,-3
                p=10,3 v=-1,2
                p=2,0 v=2,-1
                p=0,0 v=1,3
                p=3,0 v=-2,-2
                p=7,6 v=-1,-3
                p=3,0 v=-1,-2
                p=9,3 v=2,3
                p=7,3 v=-1,2
                p=2,4 v=2,-3
                p=9,5 v=-3,-3
                """
            ),
            width=11,
            height=7,
        )
        self.assertTupleEqual(
            solver.robots[0],
            Robot(0, 4, 3, -3),
        )
        self.assertTupleEqual(
            solver.move_robot(solver.robots[0], 1),
            Robot(3, 1, 3, -3),
        )
        self.assertEqual(solver.get_quadrant(Robot(0, 0, 0, 0)), 2)
        self.assertEqual(solver.get_quadrant(Robot(5, 0, 0, 0)), None)
        self.assertEqual(
            solver.total_overlap(
                [
                    Robot(0, 0, 0, 0),
                    Robot(0, 0, 0, 0),
                ]
            ),
            2,
        )
        self.assertEqual(
            solver.total_overlap(
                [
                    Robot(0, 0, 0, 0),
                    Robot(1, 0, 0, 0),
                ]
            ),
            0,
        )
        self.assertEqual(solver.solve_p1(), 12)