from tools.aoc import AOCDay
from tools.coordinate import Coordinate, DistanceAlgorithm
from tools.grid import Grid
from typing import Any


def expand_universe(grid: Grid, amount: int = 2):
    """amount as in 'one empty row should be {amount} empty rows instead"""
    empty_x = set(grid.rangeX())
    empty_y = set(grid.rangeY())

    for c in grid.getActiveCells():
        if c.x in empty_x:
            empty_x.remove(c.x)

        if c.y in empty_y:
            empty_y.remove(c.y)

    for c in reversed(sorted(grid.getActiveCells())):
        shift_x = len([x for x in empty_x if x < c.x]) * (amount - 1)
        shift_y = len([y for y in empty_y if y < c.y]) * (amount - 1)
        grid.set(c, False)
        grid.set(Coordinate(c.x + shift_x, c.y + shift_y), True)


def get_sum_dist(grid) -> int:
    dist_sum = 0
    c_list = list(grid.getActiveCells())
    for i, a in enumerate(c_list[:-1]):
        for b in c_list[i + 1 :]:
            dist_sum += a.getDistanceTo(b, algorithm=DistanceAlgorithm.MANHATTAN)

    return dist_sum


class Day(AOCDay):
    inputs = [
        [
            (374, "input11_test"),
            (9686930, "input11_dennis"),
            (9445168, "input11"),
        ],
        [
            (630728425490, "input11_dennis"),
            (742305960572, "input11"),
        ],
    ]

    def part1(self) -> Any:
        grid = Grid.from_data(self.getInput(), translate={".": False, "#": True})
        expand_universe(grid)
        return get_sum_dist(grid)

    def part2(self) -> Any:
        grid = Grid.from_data(self.getInput(), translate={".": False, "#": True})
        expand_universe(grid, 1_000_000)
        return get_sum_dist(grid)


if __name__ == "__main__":
    day = Day(2023, 11)
    day.run(verbose=True)