import re
from textwrap import dedent
from typing import Iterator, NamedTuple, override
from unittest import TestCase

from puzzles._solver import Solver


class ClawMachine(NamedTuple):
    ax: int
    ay: int
    bx: int
    by: int
    px: int
    py: int


a_spec_pattern = re.compile(r"Button A: X\+(\d+), Y\+(\d+)")
b_spec_pattern = re.compile(r"Button B: X\+(\d+), Y\+(\d+)")
p_spec_pattern = re.compile(r"Prize: X=(\d+), Y=(\d+)")


def solve_machine(m: ClawMachine) -> tuple[float, float]:
    c = (m.px * m.by - m.py * m.bx) / (m.ax * m.by - m.ay * m.bx)
    d = (m.px - m.ax * c) / m.bx
    return c, d


class DayThirteenSolver(Solver):
    machines: list[ClawMachine]

    @override
    def __init__(self, puzzle_input: str):
        self.machines = []
        for spec in puzzle_input.strip().split("\n\n"):
            a_spec, b_spec, p_spec = spec.split("\n")
            ax, ay = map(int, a_spec_pattern.match(a_spec).groups())
            bx, by = map(int, b_spec_pattern.match(b_spec).groups())
            px, py = map(int, p_spec_pattern.match(p_spec).groups())
            self.machines.append(ClawMachine(ax, ay, bx, by, px, py))

    @override
    def solve_p1(self) -> int:
        return sum(
            int(c * 3 + d)
            for c, d in map(solve_machine, self.machines)
            if c.is_integer() and d.is_integer()
        )

    @override
    def solve_p2(self) -> int:
        corrected_machines = map(
            lambda m: ClawMachine(
                m.ax,
                m.ay,
                m.bx,
                m.by,
                m.px + 10000000000000,
                m.py + 10000000000000,
            ),
            self.machines,
        )
        return sum(
            int(c * 3 + d)
            for c, d in map(solve_machine, corrected_machines)
            if c.is_integer() and d.is_integer()
        )


class TestDayThirteenSolver(TestCase):
    def test(self):
        solver = DayThirteenSolver(
            dedent(
                """
                Button A: X+94, Y+34
                Button B: X+22, Y+67
                Prize: X=8400, Y=5400

                Button A: X+26, Y+66
                Button B: X+67, Y+21
                Prize: X=12748, Y=12176

                Button A: X+17, Y+86
                Button B: X+84, Y+37
                Prize: X=7870, Y=6450

                Button A: X+69, Y+23
                Button B: X+27, Y+71
                Prize: X=18641, Y=10279
                """
            )
        )
        self.assertEqual(
            solver.machines[0],
            ClawMachine(94, 34, 22, 67, 8400, 5400),
        )
        self.assertEqual(solver.solve_p1(), 480)