import numpy as np
from pathlib import Path


def pairs_from_list(array):
    if len(array) == 2:
        return [tuple(array)]
    return [(array[0], k) for k in array[1:]] + pairs_from_list(array[1:])


class AntennaMap:
    def __init__(self, filepath):
        with open(filepath, "r") as filein:
            self.map = np.array(
                [[ord(c) for c in line.rstrip()] for line in filein.readlines()]
            )
            self.xlen, self.ylen = self.map.shape
            unifrq = np.unique(self.map)
            unique_freqs = np.delete(unifrq, np.argwhere(unifrq == ord(".")))
            self.nodes = set()
            self.antennae = {k: np.argwhere(self.map == k) for k in unique_freqs}

    def __str__(self):
        map = np.copy(self.map)
        for i, j in self.nodes:
            map[i, j] = ord("#") if map[i, j] == ord(".") else map[i, j]
        return "\n".join(["".join([chr(c) for c in line]) for line in map])

    def _on_map(self, coords):
        x, y = coords
        return (0 <= x < self.xlen) and (0 <= y < self.ylen)

    def find_antinode_pairs(self, freq):
        coords = self.antennae[freq]
        pairs = pairs_from_list(coords)
        for x1, x2 in pairs:
            u = np.subtract(x2, x1)
            x3 = x2 + u
            x0 = x1 - u
            if self._on_map(x0):
                self.nodes.add(tuple(x0))
            if self._on_map(x3):
                self.nodes.add(tuple(x3))

    def _border_distance_2d(self, coords, vect):
        result = [
            self._border_distance_1d(x, u, mx)
            for x, u, mx in zip(coords, vect, self.map.shape)
        ]
        return np.min(result)

    def _border_distance_1d(self, x, u, mx):
        xb = 0 if u < 0 else mx - 1
        length = xb - x
        return length // u

    def find_antinodes(self, freq):
        coords = self.antennae[freq]
        pairs = pairs_from_list(coords)
        for x1, x2 in pairs:
            u = np.subtract(x2, x1)
            k0 = self._border_distance_2d(x2, -u)
            k1 = self._border_distance_2d(x2, u)
            for k in range(-k0, k1 + 1):
                self.nodes.add(tuple(x2 + k * u))

    def scan_prime(self):
        for freq in self.antennae:
            self.find_antinode_pairs(freq)

    def scan_harmonics(self):
        for freq in self.antennae:
            self.find_antinodes(freq)


if __name__ == "__main__":
    filepath = Path("./example")
    filepath = Path("./input")

    map = AntennaMap(filepath)
    map.scan_prime()
    print(map)
    print(len(map.nodes))

    map.scan_harmonics()
    print(map)
    print(len(map.nodes))