from typing import Any, Iterator

from ..base import BaseSolver


def reachable(
    map: list[str], tiles: set[tuple[int, int]], steps: int
) -> set[tuple[int, int]]:
    n_rows, n_cols = len(map), len(map[0])

    for _ in range(steps):
        tiles = {
            (i + di, j + dj)
            for i, j in tiles
            for di, dj in ((-1, 0), (+1, 0), (0, -1), (0, +1))
            if map[(i + di) % n_rows][(j + dj) % n_cols] != "#"
        }
    return tiles


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        map = input.splitlines()
        start = next(
            (i, j)
            for i in range(len(map))
            for j in range(len(map[i]))
            if map[i][j] == "S"
        )

        # part 1
        yield len(reachable(map, {start}, 6 if len(map) < 20 else 64))

        # part 2

        # the initial map is a square and contains an empty rhombus whose diameter is
        # the size of the map, and has only empty cells around the middle row and column
        #
        # after ~n/2 steps, the first map is filled with a rhombus, after that we get a
        # bigger rhombus every n steps
        #
        # we are going to find the number of cells reached for the initial rhombus, n
        # steps after and n * 2 steps after
        #
        cycle = len(map)
        rhombus = (len(map) - 3) // 2 + 1

        values: list[int] = []
        values.append(len(tiles := reachable(map, {start}, rhombus)))
        values.append(len(tiles := reachable(map, tiles, cycle)))
        values.append(len(tiles := reachable(map, tiles, cycle)))

        if self.verbose:
            n_rows, n_cols = len(map), len(map[0])

            rows = [
                [
                    map[i % n_rows][j % n_cols] if (i, j) not in tiles else "O"
                    for j in range(-2 * cycle, 3 * cycle)
                ]
                for i in range(-2 * cycle, 3 * cycle)
            ]

            for i in range(len(rows)):
                for j in range(len(rows[i])):
                    if (i // cycle) % 2 == (j // cycle) % 2:
                        rows[i][j] = f"\033[91m{rows[i][j]}\033[0m"

            for row in rows:
                self.logger.info("".join(row))

        self.logger.info(f"values to fit: {values}")

        # version 1:
        #
        # after 3 cycles, the figure looks like the following:
        #
        #   I M D
        # I J A K D
        # H A F A L
        # C E A K B
        #   C G B
        #
        # after 4 cycles, the figure looks like the following:
        #
        #     I M D
        #   I J A K D
        # I J A B A K D
        # H A B A B A L
        # C E A B A N F
        #   C E A N F
        #     C G F
        #
        # the 'radius' of the rhombus is the number of cycles minus 1
        #
        # the 4 'corner' (M, H, L, G) are counted once, the blocks with a corner triangle (D, I,
        # C, B) are each counted radius times, the blocks with everything but one corner (J, K,
        # E, N) are each counted radius - 1 times
        #
        # there are two versions of the whole block, A and B in the above (or odd and even),
        # depending on the number of cycles, either A or B will be in the center
        #

        counts = [
            [
                sum(
                    (i, j) in tiles
                    for i in range(ci * cycle, (ci + 1) * cycle)
                    for j in range(cj * cycle, (cj + 1) * cycle)
                )
                for cj in range(-2, 3)
            ]
            for ci in range(-2, 3)
        ]

        radius = (26501365 - rhombus) // cycle - 1
        A = counts[2][2] if radius % 2 == 0 else counts[2][1]
        B = counts[2][2] if radius % 2 == 1 else counts[2][1]
        answer_2 = (
            (radius + 1) * A
            + radius * B
            + 2 * radius * (radius + 1) // 2 * A
            + 2 * radius * (radius - 1) // 2 * B
            + sum(counts[i][j] for i, j in ((0, 2), (-1, 2), (2, 0), (2, -1)))
            + sum(counts[i][j] for i, j in ((0, 1), (0, 3), (-1, 1), (-1, 3)))
            * (radius + 1)
            + sum(counts[i][j] for i, j in ((1, 1), (1, 3), (-2, 1), (-2, 3))) * radius
        )
        print(f"answer 2 (v1) is {answer_2}")

        # version 2: fitting a polynomial
        #
        # the value we are interested in (26501365) can be written as R + K * C where R is the
        # step at which we find the first rhombus, and K the repeat step, so instead of fitting
        # for X values (R, R + K, R + 2 K), we are going to fit for (0, 1, 2), giving us much
        # simpler equation for the a, b and c coefficient
        #
        # we get:
        #   - (a * 0² + b * 0 + c) = y1  =>  c = y1
        #   - (a * 1² + b * 1 + c) = y2  =>  a + b = y2 - y1
        #                                =>  b = y2 - y1 - a
        #   - (a * 2² + b * 2 + c) = y3  =>  4a + 2b = y3 - y1
        #                                =>  4a + 2(y2 - y1 - a) = y3 - y1
        #                                =>  a = (y1 + y3) / 2 - y2
        #
        y1, y2, y3 = values
        a, b, c = (y1 + y3) // 2 - y2, 2 * y2 - (3 * y1 + y3) // 2, y1

        n = (26501365 - rhombus) // cycle
        yield a * n * n + b * n + c