import itertools as it
from typing import Any, Iterator

import numpy as np

from ..base import BaseSolver


def move(head: tuple[int, int], command: str) -> tuple[int, int]:
    h_col, h_row = head

    if command == "L":
        head = (h_col - 1, h_row)
    elif command == "R":
        head = (h_col + 1, h_row)
    elif command == "U":
        head = (h_col, h_row + 1)
    elif command == "D":
        head = (h_col, h_row - 1)

    return head


def follow(head: tuple[int, int], tail: tuple[int, int]) -> tuple[int, int]:
    h_col, h_row = head
    t_col, t_row = tail

    if abs(t_col - h_col) <= 1 and abs(t_row - h_row) <= 1:
        return tail

    return t_col + np.sign(h_col - t_col), t_row + np.sign(h_row - t_row)


def run(commands: list[str], n_blocks: int) -> list[tuple[int, int]]:
    blocks: list[tuple[int, int]] = [(0, 0) for _ in range(n_blocks)]
    visited = [blocks[-1]]

    for command in commands:
        blocks[0] = move(blocks[0], command)

        for i in range(0, n_blocks - 1):
            blocks[i + 1] = follow(blocks[i], blocks[i + 1])

        visited.append(blocks[-1])

    return visited


class Solver(BaseSolver):
    def solve(self, input: str) -> Iterator[Any]:
        lines = [line.strip() for line in input.splitlines()]

        # flatten the commands
        commands = list(
            it.chain(*(p[0] * int(p[1]) for line in lines if (p := line.split())))
        )

        yield len(set(run(commands, n_blocks=2)))
        yield len(set(run(commands, n_blocks=10)))