from typing import Any, Iterator, Sequence, TypeAlias, TypeVar

import numpy as np
from numpy.typing import NDArray

from ..base import BaseSolver

T = TypeVar("T")

Tower: TypeAlias = NDArray[np.bool]


def tower_height(tower: Tower) -> int:
    return int(tower.shape[0] - tower[::-1, :].argmax(axis=0).min() - 1)


def next_cycle(sequence: Sequence[T], index: int) -> tuple[T, int]:
    t = sequence[index]
    index = (index + 1) % len(sequence)
    return t, index


ROCKS = [
    np.array([(0, 0), (0, 1), (0, 2), (0, 3)]),
    np.array([(0, 1), (1, 0), (1, 1), (1, 2), (2, 1)]),
    np.array([(0, 0), (0, 1), (0, 2), (1, 2), (2, 2)]),
    np.array([(0, 0), (1, 0), (2, 0), (3, 0)]),
    np.array([(0, 0), (0, 1), (1, 0), (1, 1)]),
]

WIDTH = 7
START_X = 2

EMPTY_BLOCKS = np.zeros((10, WIDTH), dtype=bool)


def build_tower(
    n_rocks: int,
    jets: str,
    early_stop: bool = False,
    init: Tower = np.ones(WIDTH, dtype=bool),
) -> tuple[Tower, int, int, dict[int, int]]:
    tower = EMPTY_BLOCKS.copy()
    tower[0, :] = init

    done_at: dict[tuple[int, int], int] = {}
    heights: dict[int, int] = {}
    i_jet, i_rock = 0, 0
    rock_count = 0

    for rock_count in range(n_rocks):
        if early_stop:
            if i_rock == 0 and (i_rock, i_jet) in done_at:
                break
            done_at[i_rock, i_jet] = rock_count

        y_start = tower.shape[0] - tower[::-1, :].argmax(axis=0).min() + 3
        rock, i_rock = next_cycle(ROCKS, i_rock)

        rock_y = rock[:, 0] + y_start
        rock_x = rock[:, 1] + START_X

        if rock_y.max() >= tower.shape[0]:
            tower = np.concatenate([tower, EMPTY_BLOCKS], axis=0)

        while True:
            jet, i_jet = next_cycle(jets, i_jet)

            dx = 0
            if jet == ">" and rock_x.max() < WIDTH - 1:
                dx = 1
            elif jet == "<" and rock_x.min() > 0:
                dx = -1

            if dx != 0 and not tower[rock_y, rock_x + dx].any():
                rock_x = rock_x + dx

            # move down
            rock_y -= 1

            if tower[rock_y, rock_x].any():
                rock_y += 1
                break

        heights[rock_count] = tower_height(tower)
        tower[rock_y, rock_x] = True

    return tower, rock_count, done_at.get((i_rock, i_jet), -1), heights


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        tower, *_ = build_tower(2022, input)
        yield tower_height(tower)

        TOTAL_ROCKS = 1_000_000_000_000
        _tower_1, n_rocks_1, prev_1, heights_1 = build_tower(TOTAL_ROCKS, input, True)
        assert prev_1 > 0

        # 2767 1513
        remaining_rocks = TOTAL_ROCKS - n_rocks_1
        n_repeat_rocks = n_rocks_1 - prev_1
        n_repeat_towers = remaining_rocks // n_repeat_rocks

        base_height = heights_1[prev_1]
        repeat_height = heights_1[prev_1 + n_repeat_rocks - 1] - heights_1[prev_1]
        remaining_height = (
            heights_1[prev_1 + remaining_rocks % n_repeat_rocks] - heights_1[prev_1]
        )

        yield base_height + (n_repeat_towers + 1) * repeat_height + remaining_height