import logging
import os
import sys

VERBOSE = os.getenv("AOC_VERBOSE") == "True"
logging.basicConfig(level=logging.INFO if VERBOSE else logging.WARNING)


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


map = sys.stdin.read().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
answer_1 = len(reachable(map, {start}, 6 if len(map) < 20 else 64))
print(f"answer 1 is {answer_1}")

# 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, and then interpolate the value to get a 2nd order
# polynomial
#
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 logging.root.getEffectiveLevel() == logging.INFO:
    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 + 1)
        ]
        for i in range(-2 * cycle, 3 * cycle + 1)
    ]

    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"

    print("\n".join("".join(row) for row in rows))


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

# 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
answer_2 = a * n * n + b * n + c
print(f"answer 2 is {answer_2}")