from itertools import combinations
from auxiliares import (
    string_para_matriz,
    Vector2d,
    Labirinto,
    le_arquivo_sem_linhas_em_branco,
)

ENTRADA = """
............
........0...
.....0......
.......0....
....0.......
......A.....
............
............
........A...
.........A..
............
............
"""


class Campo(Labirinto):
    def __init__(self, labirinto: list[list[str]]):
        super().__init__(labirinto)
        self.anti_sinais = {}
        self.acha_antenas()

    def desenha(self, percurso: dict[str, list[Vector2d]], caractere="X"):
        for percurso in percurso.values():
            for posição in percurso:
                if self.posição_válida(posição):
                    self[posição.y, posição.x] = caractere

    def acha_antenas(self):
        self.antenas = {}
        for y, linha in enumerate(self.labirinto):
            for x, pos in enumerate(linha):
                if pos != ".":
                    self.antenas[pos] = self.antenas.get(pos, []) + [Vector2d(y, x)]

    def gera_anti_sinal(self):
        self.anti_sinais = {}
        for antena, posições in self.antenas.items():
            if len(posições) > 1:
                for p1, p2 in combinations(posições, 2):
                    d = p2 - p1
                    if self.posição_válida(as1 := p1 - d):
                        self.novo_anti_sinal(antena, as1)
                    if self.posição_válida(as2 := p2 + d):
                        self.novo_anti_sinal(antena, as2)

    def novo_anti_sinal(self, antena, posição: Vector2d):
        self.anti_sinais[antena] = self.anti_sinais.get(antena, []) + [posição]

    def gera_anti_sinal_com_resonância(self):
        self.anti_sinais = {}
        for antena, posições in self.antenas.items():
            for posição in self.antenas[antena]:
                self.novo_anti_sinal(antena, posição)
            if len(posições) > 1:
                for p1, p2 in combinations(posições, 2):
                    d = p2 - p1
                    while self.posição_válida(as1 := p1 - d):
                        self.novo_anti_sinal(antena, as1)
                        p1 = as1
                    while self.posição_válida(as2 := p2 + d):
                        self.novo_anti_sinal(antena, as2)
                        p2 = as2

    def conta_anti_sinais(self):
        únicos = set()
        for a in self.anti_sinais.values():
            for anti_sinal in a:
                únicos.add(anti_sinal)
        return len(únicos)


teste_entrada = string_para_matriz(ENTRADA)
campo = Campo(teste_entrada)
campo.gera_anti_sinal()
resposta_parte_1_teste = campo.conta_anti_sinais()
print("anti-sinais únicos (teste - parte 1):", resposta_parte_1_teste)
assert resposta_parte_1_teste == 14

entrada_parte_1 = string_para_matriz(le_arquivo_sem_linhas_em_branco("08/input.txt"))
campo_parte_1 = Campo(entrada_parte_1)
campo_parte_1.gera_anti_sinal()
resposta_parte_1 = campo_parte_1.conta_anti_sinais()
print("anti-sinais únicos (parte 1):", resposta_parte_1)
assert resposta_parte_1 == 276

campo = Campo(teste_entrada)
campo.gera_anti_sinal_com_resonância()
resposta_parte2_teste = campo.conta_anti_sinais()
print("anti-sinais únicos com resonância (teste - parte 2):", resposta_parte2_teste)
assert resposta_parte2_teste == 34

campo_parte_2 = Campo(entrada_parte_1)
campo_parte_2.gera_anti_sinal_com_resonância()
resposta_parte_2 = campo_parte_2.conta_anti_sinais()
print("anti-sinais únicos com resonância (parte 2):", resposta_parte_2)
assert resposta_parte_2 == 991

# Saída
# anti-sinais únicos (teste - parte 1): 14
# anti-sinais únicos (parte 1): 276
# anti-sinais únicos com resonância (teste - parte 2): 34
# anti-sinais únicos com resonância (parte 2): 991

class Vector2d:
    def __init__(self, y, x):
        self.x = x
        self.y = y

    def copy(self):
        return Vector2d(self.y, self.x)

    def __add__(self, other):
        return Vector2d(self.y + other.y, self.x + other.x)

    def __sub__(self, other):
        return Vector2d(self.y - other.y, self.x - other.x)

    def __str__(self):
        return f"({self.y}, {self.x})"

    def __repr__(self):
        return f"Vector2d({self.y}, {self.x})"

    def __hash__(self):
        return hash(f"({self.y}, {self.x})")

    def __eq__(self, other):
        if isinstance(other, Vector2d):
            return self.x == other.x and self.y == other.y
        return False


class Labirinto:
    def __init__(self, labirinto):
        self.max_y = len(labirinto)
        self.max_x = len(labirinto[0])
        self.orig_labirinto = labirinto
        self.labirinto = [list(linha) for linha in labirinto]

    def copy(self):
        return Labirinto(self.orig_labirinto)

    def posição_válida(self, posição: Vector2d) -> bool:
        return 0 <= posição.y < self.max_y and 0 <= posição.x < self.max_x

    def __str__(self):
        return "\n".join(["".join(linha) for linha in self.labirinto])

    def __setitem__(self, yx, value):
        self.labirinto[yx[0]][yx[1]] = value

    def __getitem__(self, yx):
        return self.labirinto[yx[0]][yx[1]]


def le_arquivo_sem_linhas_em_branco(nome_arquivo: str) -> str:
    return "\n".join(
        zz
        for linha in le_arquivo(nome_arquivo).splitlines()
        if len(zz := linha.strip()) > 0
    )