from collections.abc import Generator
from itertools import combinations, product
from pathlib import Path
from typing import NamedTuple


class Point(NamedTuple):
    x: int
    y: int

    def iter_adjacent(self, wrap=False) -> Generator["Point", None, None]:
        for offset in product((1, 0, -1), (1, 0, -1)):
            adj = Point(self.x + offset[0], self.y + offset[1])
            if adj == self:
                continue
            if not wrap and (adj.x < 0 or adj.y < 0):
                continue
            yield adj

    def add_point(self, delta: "Point") -> "Point":
        return Point(self.x + delta.x, self.y + delta.y)

    def add(self, x=0, y=0) -> "Point":
        return Point(self.x + x, self.y + y)

    def dist(self, x=0, y=0, point: "Point|None" = None) -> int:
        if (x or y) and point:
            raise ValueError("Cannot provide both x and y and a point")

        if point:
            x = point.x
            y = point.y

        dx = abs(self.x - x)
        dy = abs(self.y - y)

        return dx + dy

    def __str__(self) -> str:
        return f"Point({self.x}, {self.y})"


class Galaxy:
    def __init__(self, lines: list[str]):
        self.planets: set[Point] = set()

        self.grid = [[x for x in row] for row in lines if row]
        self.expand_y: dict[int, int] = {}
        self.expand_x: dict[int, int] = {}

    def find_planets(self) -> set[Point]:
        # Find planets
        for y, row in enumerate(self.grid):
            for x, c in enumerate(row):
                if c == "#":
                    self.planets.add(Point(x, y))

        return self.planets

    def dist(self, a: Point, b: Point) -> int:
        d = a.dist(point=b)

        for y in range(*sorted((a.y, b.y))):
            d += self.expand_y.get(y, 0)

        for x in range(*sorted((a.x, b.x))):
            d += self.expand_x.get(x, 0)

        return d

    def virtual_expand(self, n=2) -> None:
        # Expand y
        self.expand_y = {
            y: n-1 for y, row in enumerate(self.grid) if all(c == "." for c in row)
        }

        # Expand x
        self.expand_x = {
            x: n-1 for x in range(len(self.grid[0])) if all(row[x] == "." for row in self.grid)

        }

    def expand(self, n = 2) -> None:
        # Expand y
        og_y = len(self.grid)
        for y in range(og_y):
            i = og_y - y - 1
            row = self.grid[i]
            if all(c == "." for c in row):
                self.grid.insert(i, [str(n-1) for _ in row])

        # Expand x
        og_x = len(self.grid[0])
        for x in range(og_x):
            i = og_x - x - 1
            col = [row[i] for row in self.grid]
            if all(c == "." or c.isnumeric() for c in col):
                for row in self.grid:
                    row.insert(i, str(n-1))

    def __str__(self) -> str:
        return "\n".join("".join(row) for row in self.grid)


def part1(input: Path) -> int:
    galaxy = Galaxy(input.read_text().split("\n"))
    galaxy.expand()

    galaxy.find_planets()

    return sum(
        a.dist(point=b) for a, b in combinations(galaxy.planets, 2)
    )


def part2(input: Path, n=1000000) -> int:
    galaxy = Galaxy(input.read_text().split("\n"))

    galaxy.virtual_expand(n)

    galaxy.find_planets()
    # print(galaxy.dist(Point(1, 0), Point(3, 1)))

    return sum(
        galaxy.dist(a, b) for a, b in combinations(galaxy.planets, 2)
    )


if __name__ == "__main__":
    # input = Path("sample.txt")
    input = Path("input.txt")

    result1 = part1(input)
    print("part 1", result1)

    print("part 2 10", part2(input, 10))
    print("part 2 100", part2(input, 100))
    print("part 2 1M", part2(input))