from collections import defaultdict
from itertools import combinations
from tools.aoc import AOCDay
from tools.coordinate import Coordinate, DistanceAlgorithm
from tools.grid import Grid, GridTransformation
from typing import Any, List, Union


ROTATION_PATH = [
    GridTransformation.ROTATE_X,
    GridTransformation.ROTATE_X,
    GridTransformation.ROTATE_Z,
    GridTransformation.ROTATE_X,
    GridTransformation.ROTATE_X,
    GridTransformation.ROTATE_Z
]


def overlap(beacon1: Grid, beacon2: Grid) -> Union[Coordinate, None]:
    diffs = defaultdict(int)
    for c1 in beacon1.getActiveCells():
        for c2 in beacon2.getActiveCells():
            # Interestingly adding an if statement here to break out early when already 12 matches are found
            # slows the loop down so much, that the end result is slower compared to just counting everything
            # and checking the results afterwards.
            diffs[c1 - c2] += 1

    for d in diffs:
        if diffs[d] >= 12:
            return d


def converge(scanner_grids: List[Grid]) -> (Grid, List[Coordinate]):
    main_grid = scanner_grids.pop()
    diff_list = [Coordinate(0, 0, 0)]
    while scanner_grids:
        found = False
        for i, scanner in enumerate(scanner_grids):
            for rotation in ROTATION_PATH:
                for y_rotation in range(4):
                    if diff := overlap(main_grid, scanner):
                        #print("Grid overlap:", main_grid.name, scanner.name, "=>", len(scanner_grids) - 1, "to go.")
                        scanner_grids.pop(i)
                        diff_list.append(diff)
                        for c in scanner.getActiveCells():
                            main_grid.set(c + diff)
                        found = True
                        break
                    scanner.transform(GridTransformation.ROTATE_Y)
                scanner.transform(rotation)
                if found:
                    break
            if found:
                break
        assert found, "No overlapping Grids found"

    return main_grid, diff_list


class Day(AOCDay):
    inputs = [
        [
            (79, "test_input19"),
            (465, "input19")
        ],
        [
            (3621, "test_input19"),
            (12149, "input19")
        ]
    ]

    def getScannerGrids(self) -> List[Grid]:
        scanner_grids = []
        for scanner_input in self.getMultiLineInputAsArray():
            scanner_grid = Grid()
            scanner_grid.name = scanner_input[0]
            for l in scanner_input[1:]:
                scanner_grid.set(Coordinate(*map(int, l.split(","))))
            scanner_grids.append(scanner_grid)

        return scanner_grids

    def part1(self) -> Any:
        teh_grid, _ = converge(self.getScannerGrids())

        return len(teh_grid.getActiveCells())

    def part2(self) -> Any:
        _, diff_list = converge(self.getScannerGrids())
        max_diff = 0
        for s1, s2 in combinations(diff_list, 2):
            max_diff = max(max_diff, s1.getDistanceTo(s2, includeDiagonals=False, algorithm=DistanceAlgorithm.MANHATTAN))

        return max_diff


if __name__ == '__main__':
    day = Day(2021, 19)
    day.run(verbose=True)