from __future__ import annotations

from array import array
from collections.abc import Iterable
from dataclasses import dataclass

from .aoc import Aoc2, AocSameParser


@dataclass
class Position:
    x: int
    y: int

    def __add__(self, other: Position) -> Position:
        return Position(self.x + other.x, self.y + other.y)

    def __sub__(self, other: Position) -> Position:
        return Position(self.x - other.x, self.y - other.y)

    def __or__(self, other: Position) -> int:
        return abs(self.x - other.x) + abs(self.y - other.y)

    def __rmul__(self, other: int) -> Position:
        return Position(other * self.x, other * self.y)

    def __getitem__(self, pi: int) -> int:
        if pi == 0:
            return self.x
        return self.y

    def __repr__(self) -> str:
        return f"[{self.x},{self.y}]"


class Matrix:
    def __init__(self, inpt: str) -> None:
        lines = inpt.splitlines()
        self.height = len(lines)
        self.width = len(lines[0])
        self.array = array("u")
        for line in lines:
            self.array.fromlist(list(line))

    def __getitem__(self, pos: Position | tuple[int, int]) -> str:
        return self.array[self.width * pos[1] + pos[0]]

    def __setitem__(self, pos: Position | tuple[int, int], value: str) -> None:
        self.array[self.width * pos[1] + pos[0]] = value

    def get_row(self, row_no: int) -> array[str]:
        return self.array[self.width * row_no : self.width * (row_no + 1)]

    def iter_rows(self) -> Iterable[array[str]]:
        return (self.get_row(y) for y in range(self.height))

    def get_col(self, col_no: int) -> array[str]:
        return array("u", [self[col_no, y] for y in range(self.height)])

    def iter_cols(self) -> Iterable[array[str]]:
        return (self.get_col(x) for x in range(self.width))

    def __str__(self) -> str:
        return "\n".join((x.tounicode() for x in self.iter_rows()))

    def finditer(self, element: str) -> Iterable[Position]:
        for y, line in enumerate(self.iter_rows()):
            for x, char in enumerate(line):
                if char == element:
                    yield Position(x, y)


@dataclass
class Map:
    map: Matrix
    galaxies: list[Position]
    adjusted_galaxies: list[Position]
    large_adjusted_galaxies: list[Position]

    @classmethod
    def from_input(cls, inpt: str) -> Map:
        matrix = Matrix(inpt)

        galaxies = list(matrix.finditer("#"))
        missing_ys = [
            y for y, row in enumerate(matrix.iter_rows()) if row.count("#") == 0
        ]
        missing_xs = [
            x for x, col in enumerate(matrix.iter_cols()) if col.count("#") == 0
        ]

        adjusted_galaxies = []
        large_adjusted_galaxies = []

        for pos in galaxies:
            offset = Position(
                len([m for m in missing_xs if m < pos.x]),
                len([m for m in missing_ys if m < pos.y]),
            )
            adjusted_galaxies.append(pos + offset)
            large_adjusted_galaxies.append(pos + 999_999 * offset)

        return cls(matrix, galaxies, adjusted_galaxies, large_adjusted_galaxies)


class Day11(AocSameParser[Map], Aoc2[Map, Map]):
    def parseinput(self, inpt: str) -> Map:
        return Map.from_input(inpt)

    def part1(self, inpt: Map) -> int:
        return sum(
            g1 | g2
            for i, g1 in enumerate(inpt.adjusted_galaxies)
            for g2 in inpt.adjusted_galaxies[i + 1 :]
        )

    def part2(self, inpt: Map) -> int:
        return sum(
            g1 | g2
            for i, g1 in enumerate(inpt.large_adjusted_galaxies)
            for g2 in inpt.large_adjusted_galaxies[i + 1 :]
        )


def main() -> None:
    day11 = Day11(2023, 11)
    day11.run()


if __name__ == "__main__":
    main()