from collections import defaultdict

from tools.aoc import AOCDay
from typing import Any

STR_OP_CODES = [
    "addr",
    "addi",
    "mulr",
    "muli",
    "banr",
    "bani",
    "borr",
    "bori",
    "setr",
    "seti",
    "gtir",
    "gtri",
    "gtrr",
    "eqir",
    "eqri",
    "eqrr",
]


class WristDevice:
    def __init__(self):
        self.registers = [0, 0, 0, 0]
        self.op_code_learner: dict[int, set[str]] = defaultdict(set)
        self.op_code_by_id: dict[int, str] = {}

    def set_reg(self, a: int, b: int, c: int, d: int) -> None:
        self.registers = [a, b, c, d]

    def execute_by_name(self, op_code_name: str, arg0: int, arg1: int, arg2: int):
        if op_code_name == "addr":
            self.registers[arg2] = self.registers[arg0] + self.registers[arg1]
        elif op_code_name == "addi":
            self.registers[arg2] = self.registers[arg0] + arg1
        elif op_code_name == "mulr":
            self.registers[arg2] = self.registers[arg0] * self.registers[arg1]
        elif op_code_name == "muli":
            self.registers[arg2] = self.registers[arg0] * arg1
        elif op_code_name == "banr":
            self.registers[arg2] = self.registers[arg0] & self.registers[arg1]
        elif op_code_name == "bani":
            self.registers[arg2] = self.registers[arg0] & arg1
        elif op_code_name == "borr":
            self.registers[arg2] = self.registers[arg0] | self.registers[arg1]
        elif op_code_name == "bori":
            self.registers[arg2] = self.registers[arg0] | arg1
        elif op_code_name == "setr":
            self.registers[arg2] = self.registers[arg0]
        elif op_code_name == "seti":
            self.registers[arg2] = arg0
        elif op_code_name == "gtir":
            self.registers[arg2] = int(arg0 > self.registers[arg1])
        elif op_code_name == "gtri":
            self.registers[arg2] = int(self.registers[arg0] > arg1)
        elif op_code_name == "gtrr":
            self.registers[arg2] = int(self.registers[arg0] > self.registers[arg1])
        elif op_code_name == "eqir":
            self.registers[arg2] = int(arg0 == self.registers[arg1])
        elif op_code_name == "eqri":
            self.registers[arg2] = int(self.registers[arg0] == arg1)
        elif op_code_name == "eqrr":
            self.registers[arg2] = int(self.registers[arg0] == self.registers[arg1])

    def execute_by_id(self, op_code: int, arg0: int, arg1: int, arg2: int) -> None:
        if op_code not in self.op_code_by_id:
            raise ValueError(f"Unknown OP Code: {op_code}")

        self.execute_by_name(self.op_code_by_id[op_code], arg0, arg1, arg2)

    def learn(self, op_code: int, possibles: set[str]) -> None:
        if not self.op_code_learner[op_code]:
            self.op_code_learner[op_code] = possibles
        else:
            self.op_code_learner[op_code] &= possibles

    def finish_learn(self) -> None:
        while len(self.op_code_by_id) < len(self.op_code_learner):
            for op_code, possibles in self.op_code_learner.items():
                if len(possibles) == 0:
                    continue

                if len(possibles) == 1:
                    self.op_code_by_id[op_code] = possibles.pop()
                    for remove_code in self.op_code_learner:
                        if self.op_code_by_id[op_code] in self.op_code_learner[remove_code]:
                            self.op_code_learner[remove_code].remove(self.op_code_by_id[op_code])

    def test_code(self, op_code: list[int], reg_before: list[int], reg_after: list[int]) -> set[str]:
        possible_codes = set()
        for test_code in STR_OP_CODES:
            self.set_reg(*reg_before)
            self.execute_by_name(test_code, *op_code[1:])
            if self.registers == reg_after:
                possible_codes.add(test_code)

        return possible_codes


class Day(AOCDay):
    inputs = [
        [
            (1, "input16_test1"),
            (580, "input16"),
        ],
        [
            (537, "input16"),
        ],
    ]

    def parse_input(self) -> (list[(list[int], list[int], list[int])], list[list[int]]):
        """
        first return value is a list of (op_code, reg_before, reg_after)
        second one is just a list of op_code
        """
        lines = self.getInput()
        index = 0

        tests = []
        op_codes = []

        while index < len(lines):
            if lines[index].startswith("Before:"):
                reg_before = list(map(int, lines[index][:-1].split("[")[1].split(", ")))
                op_code = list(map(int, lines[index + 1].split()))
                reg_after = list(map(int, lines[index + 2][:-1].split("[")[1].split(", ")))
                tests.append((op_code, reg_before, reg_after))
                index += 4
            elif not lines[index]:
                index += 1
            else:
                op_codes.append(list(map(int, lines[index].split())))
                index += 1

        return tests, op_codes

    def part1(self) -> Any:
        tests, _ = self.parse_input()
        count = 0
        device = WristDevice()

        for test in tests:
            if len(device.test_code(*test)) >= 3:
                count += 1

        return count

    def part2(self) -> Any:
        tests, opcodes = self.parse_input()
        device = WristDevice()

        for test in tests:
            device.learn(test[0][0], device.test_code(*test))

        device.finish_learn()
        device.set_reg(0, 0, 0, 0)
        for opcode in opcodes:
            device.execute_by_id(*opcode)

        return device.registers[0]


if __name__ == "__main__":
    day = Day(2018, 16)
    day.run(verbose=True)