
# Python imports
from itertools import combinations
from pathlib import Path
from typing import Callable, Dict, Generator, Iterable, List, Optional, Set, Tuple

import networkx as nx
import numpy as np

from scipy.spatial.transform import Rotation

# Paths to data
testpath = Path("day19_test.txt")
datapath = Path("day19_data.txt")


class Scanner:
    
    """Represents a scanner, and its observable beacons"""
    
    def __init__(self, scanner_id: int, beacons: List[List[int]]) -> None:
        """Instantiate Scanner object
        
        :param scanner_id:  scanner ID number
        :param beacons:  list of beacon co-ordinates relative to scanner
        
        Scanner is initiated with location (0, 0, 0)
        """
        self.scanner_id = scanner_id
        self.beacons = {idx: np.array(_) for idx, _ in enumerate(beacons)}
        self.location = np.array([0, 0, 0]).astype(int)
        
        # Calculate matrix of Euclidean distances between beacons
        self.__calculate_beacon_distances()
        
    def __calculate_beacon_distances(self) -> None:
        """Calculates euclidean distance array between beacons"""
        distarray = []
        for idx1 in range(len(self.beacons)):
            dists = []
            for idx2 in range(len(self.beacons)):
                dists.append(np.linalg.norm(self.beacons[idx1] - self.beacons[idx2]))
            distarray.append(dists)
        self.beacon_distances = np.array(distarray)
        
    def intersects(self, scanner) -> bool:
        """Returns True if this scanner intersects the passed scanner
        
        :param scanner:  reference scanner to test for intersection
        
        Intersection is identified as any beacon from this scanner sharing
        at least 12 Euclidean distances with any beacon from the passed
        scanner.
        """
        for bcn1 in self.beacon_distances:
            for bcn2 in scanner.beacon_distances:
                if len(set(bcn1).intersection(set(bcn2))) >= 12:
                    return True
        return False
    
    def get_beacons(self) -> None:
        """Returns set of beacons"""
        return set(tuple(_) for _ in self.beacons.values())
    
    def get_intersecting_beacons(self, scanner) -> List:
        """Returns tuples of beacon co-ordinates that intersect
        
        :param scanner:  reference scanner to identify intersecting beacons        
        """
        beacons = []
        for row1, bcn1 in enumerate(self.beacon_distances):
            for row2, bcn2 in enumerate(scanner.beacon_distances):
                if len(set(bcn1).intersection(set(bcn2))) >= 12:
                    beacons.append((self.beacons[row1], scanner.beacons[row2]))
        return beacons

    def match_rotation(self, scanner) -> Optional[Tuple]:
        """Return rotation matrix to align intersecting beacons
        
        :param scanner:  reference scanner to test for intersection        
        """
        mobile, reference = [], []
        for bcn1, bcn2 in self.get_intersecting_beacons(scanner):
            mobile.append(bcn1)
            reference.append(bcn2)

        # Shortcut to generate all rotation matrices for the
        # orthogonal group
        drns = {idx: _ for idx, _ in enumerate(Rotation.create_group("O").as_matrix())}

        # Calculate vector of distance between first beacon in reference and all other
        # beacons in reference
        refdists = [tuple((reference[0] - _).astype(int)) for _ in reference] 
        
        # Rotate beacons until distances match
        for idx, rot in drns.items():
            rotated = [np.matmul(_, rot).astype(int) for _ in mobile]
            rotdists = [tuple((rotated[0] - _).astype(int)) for _ in rotated] 
            if rotdists == refdists:
                offset = (reference[0] - rotated[0]).astype(int)
                return rot, offset
            
        return None
        
    def reorient(self, scanner) -> None:
        """Update beacon and scanner locations with reference to scanner
        
        :param scanner:  reference scanner to align this scanner against
        """
        # Get rotation and offset wrt scanner
        rot, offset = self.match_rotation(scanner)
        
        # Rotate beacons
        self.beacons = {idx: np.matmul(_, rot).astype(int) for idx, _ in self.beacons.items()}
        
        # Offset beacons and scanner
        self.beacons = {idx: (_ + offset).astype(int) for idx, _ in self.beacons.items()}
        self.location = offset


def load_input(fpath: Path) -> List:
    """Return a collection of Scanner objects
    
    :param fpath:  Path to data file
    """
    scanners = []  # list of scanners from file

    with fpath.open("r") as ifh:        
        for line in [_.strip() for _ in ifh.readlines()]:
            if line.startswith("---"):  # new scanner header
                scanner_id = line.split()[2]
                beacons = []  # holds beacons for this scanner
            elif len(line) == 0:  # end of scanner info, instantiate
                scanners.append(Scanner(scanner_id, beacons))
            else:  # beacon co-ordinates
                beacons.append([int(_) for _ in line.split(",")])
        # Catch last scanner
        scanners.append(Scanner(scanner_id, beacons))
        
    return scanners


def reorient_scanners(scanners: List[Scanner]) -> List[Scanner]:
    """Returns all scanners, oriented with respect to first scanner in list
    
    :param scanners:  collection of Scanner objects
    """
    oriented = []  # scanners that have been placed

    # Take scanner 0 as the reference scanner
    oriented.append(scanners.pop(0))

    # Iterate over scanners until all scanners can be placed
    while len(scanners):
        scn2 = scanners.pop(0)
        for scn1 in oriented:
            if scn2.intersects(scn1):
                rot, offset = scn2.match_rotation(scn1)
                scn2.reorient(scn1)
                oriented.append(scn2)
                break
        if scn2 not in oriented:
            scanners.append(scn2)
            
    return oriented[:]

def get_unique_beacons(scanners):
    """Returns the set of unique beacon locations in the list of scanners
    
    :param scanners:  collection of oriented Scanner objects    
    """
    # Identify all unique beacons
    beacons = set()
    for scn in scanners:
        beacons = beacons.union(scn.get_beacons())
        
    return beacons


scanners = load_input(testpath)
print(f"Loaded {len(scanners)=} scanners")
scanners = reorient_scanners(scanners)
beacons = get_unique_beacons(scanners)
print(f"Number of unique beacons: {len(beacons)}")

Loaded len(scanners)=5 scanners
Number of unique beacons: 79


scanners = load_input(datapath)
print(f"Loaded {len(scanners)=} scanners")
scanners = reorient_scanners(scanners)
beacons = get_unique_beacons(scanners)
print(f"Number of unique beacons: {len(beacons)}")

Loaded len(scanners)=39 scanners
Number of unique beacons: 442


def get_max_manhattan_distance(scanners: List[Scanner]) -> int:
    """Returns maximum distance between locations of any pair of scanners
    
    :param scanners:  collection of Scanner objects    
    """
    maxdist = 0
    while len(scanners):
        scn1 = scanners.pop(0)
        for scn2 in scanners:
            maxdist = max(maxdist, np.absolute(scn1.location-scn2.location).sum())
    return maxdist


scanners = load_input(testpath)
print(f"Loaded {len(scanners)=} scanners")
scanners = reorient_scanners(scanners)
beacons = get_unique_beacons(scanners)
print(f"Number of unique beacons: {len(beacons)}")
print(f"Largest distance between scanners: {get_max_manhattan_distance(scanners)}")

Loaded len(scanners)=5 scanners
Number of unique beacons: 79
Largest distance between scanners: 3621


scanners = load_input(datapath)
print(f"Loaded {len(scanners)=} scanners")
scanners = reorient_scanners(scanners)
beacons = get_unique_beacons(scanners)
print(f"Number of unique beacons: {len(beacons)}")
print(f"Largest distance between scanners: {get_max_manhattan_distance(scanners)}")

Loaded len(scanners)=39 scanners
Number of unique beacons: 442
Largest distance between scanners: 11079

Day 19: Beacon Scanner¶

Puzzle 2¶