from __future__ import annotations
from typing import TypeAlias

from .aoc import Aoc2, AocSameParser
from .day11 import Matrix

P1: TypeAlias = Matrix
P2: TypeAlias = Matrix


class Day14(AocSameParser[P1], Aoc2[P1, P2]):
    def parseinput(self, inpt: str) -> P1:
        return Matrix(inpt)

    def part1(self, inpt: P1) -> int:
        o_pos: frozenset[tuple[int, int]] = frozenset(
            (o.x, o.y) for o in inpt.finditer("O")
        )
        for p in o_pos:
            inpt[p] = "."

        north_matrix = {}
        for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            for x in range(0, inpt.width):  # enumerate(line):
                north_matrix[x, y] = (x, y)
                if y > 0:
                    if inpt[x, y - 1] == ".":
                        north_matrix[x, y] = north_matrix[x, y - 1]
        new_o_pos = []
        for ox, oy in o_pos:
            new_x, new_y = north_matrix[ox, oy]
            while (new_x, new_y) in new_o_pos:
                new_y += 1
            new_o_pos.append((new_x, new_y))

        return sum(inpt.height - oy for _, oy in new_o_pos)

    def part2(self, inpt: P2) -> int:
        o_pos: frozenset[tuple[int, int]] = frozenset(
            (o.x, o.y) for o in inpt.finditer("O")
        )
        for p in o_pos:
            inpt[p] = "."

        north_matrix = {}
        for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            for x in range(0, inpt.width):  # enumerate(line):
                north_matrix[x, y] = (x, y)
                if y > 0:
                    if inpt[x, y - 1] == ".":
                        north_matrix[x, y] = north_matrix[x, y - 1]

        west_matrix = {}
        for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            for x in range(0, inpt.width):  # enumerate(line):
                west_matrix[x, y] = (x, y)
                if x > 0:
                    if inpt[x - 1, y] == ".":
                        west_matrix[x, y] = west_matrix[x - 1, y]

        south_matrix = {}
        for y in reversed(range(0, inpt.height)):  # enumerate(inpt.iter_rows()):
            for x in range(0, inpt.width):  # enumerate(line):
                south_matrix[x, y] = (x, y)
                if y < inpt.height - 1:
                    if inpt[x, y + 1] == ".":
                        south_matrix[x, y] = south_matrix[x, y + 1]

        east_matrix = {}
        for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            for x in reversed(range(0, inpt.width)):  # enumerate(line):
                east_matrix[x, y] = (x, y)
                if x < inpt.width - 1:
                    if inpt[x + 1, y] == ".":
                        east_matrix[x, y] = east_matrix[x + 1, y]

        cache = {}
        stepsize = -1
        loopstart = -1

        for i in range(1_000_000_000):
            # for i in range(2):
            north_pos = set()
            for ox, oy in o_pos:
                new_x, new_y = north_matrix[ox, oy]
                while (new_x, new_y) in north_pos:
                    new_y += 1
                north_pos.add((new_x, new_y))

            # west
            west_pos = set()
            for ox, oy in north_pos:
                new_x, new_y = west_matrix[ox, oy]
                while (new_x, new_y) in west_pos:
                    new_x += 1
                west_pos.add((new_x, new_y))

            # south
            south_pos = set()
            for ox, oy in west_pos:
                new_x, new_y = south_matrix[ox, oy]
                while (new_x, new_y) in south_pos:
                    new_y -= 1
                south_pos.add((new_x, new_y))

            # east
            east_pos = set()
            for ox, oy in south_pos:
                new_x, new_y = east_matrix[ox, oy]
                while (new_x, new_y) in east_pos:
                    new_x -= 1
                east_pos.add((new_x, new_y))

            o_pos = frozenset(east_pos)

            # north_matrix = {}
            # for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            #     for x in range(0, inpt.width):  # enumerate(line):
            #         north_matrix[x, y] = (x, y)
            #         if y > 0:
            #             if (x, y - 1) in o_pos:
            #                 p = north_matrix[x, y - 1]
            #                 north_matrix[x, y] = (p[0], p[1] + 1)
            #             elif inpt[x, y - 1] == ".":
            #                 north_matrix[x, y] = north_matrix[x, y - 1]
            # o_pos = frozenset(north_matrix[pos] for pos in o_pos)
            #
            # west_matrix = {}
            # for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            #     for x in range(0, inpt.width):  # enumerate(line):
            #         west_matrix[x, y] = (x, y)
            #         if x > 0:
            #             if (x - 1, y) in o_pos:
            #                 p = west_matrix[x - 1, y]
            #                 west_matrix[x, y] = (p[0] + 1, p[1])
            #             elif inpt[x - 1, y] == ".":
            #                 west_matrix[x, y] = west_matrix[x - 1, y]
            # o_pos = frozenset(west_matrix[pos] for pos in o_pos)
            #
            # south_matrix = {}
            # for y in reversed(range(0, inpt.height)):  # enumerate(inpt.iter_rows()):
            #     for x in range(0, inpt.width):  # enumerate(line):
            #         south_matrix[x, y] = (x, y)
            #         if y < inpt.height - 1:
            #             if (x, y + 1) in o_pos:
            #                 p = south_matrix[x, y + 1]
            #                 south_matrix[x, y] = (p[0], p[1] - 1)
            #             elif inpt[x, y + 1] == ".":
            #                 south_matrix[x, y] = south_matrix[x, y + 1]
            # o_pos = frozenset(south_matrix[pos] for pos in o_pos)
            #
            # east_matrix = {}
            # for y in range(0, inpt.height):  # enumerate(inpt.iter_rows()):
            #     for x in reversed(range(0, inpt.width)):  # enumerate(line):
            #         east_matrix[x, y] = (x, y)
            #         if x < inpt.width - 1:
            #             if (x + 1, y) in o_pos:
            #                 p = east_matrix[x + 1, y]
            #                 east_matrix[x, y] = (p[0] - 1, p[1])
            #             elif inpt[x + 1, y] == ".":
            #                 east_matrix[x, y] = east_matrix[x + 1, y]
            # o_pos = frozenset(east_matrix[pos] for pos in o_pos)
            if o_pos in cache:
                if stepsize < 0:
                    loopstart = cache[o_pos]
                    stepsize = i - cache[o_pos]
                if cache[o_pos] == (1000000000 - loopstart) % stepsize + loopstart - 1:
                    break
            else:
                cache[o_pos] = i

        accum = 0
        for _, oy in o_pos:
            accum += inpt.height - oy
        for p in o_pos:
            inpt[p] = "O"
        return accum


def main() -> None:
    aoc = Day14(2023, 14, example_code_nr2=0)
    aoc.run()


if __name__ == "__main__":
    main()