"""
3D-шахматы (Raumschach-стиль) — однопроходной генератор ходов.

Доска : N×N×N  (по умолчанию 8×8×8)
Координата : (x, y, z) — колонка, ряд, уровень

Фигуры и направления движения в трёхмерном пространстве:
──────────────────────────────────────────────────────────
  Ладья    (R) — 6 осевых направлений
  Слон     (B) — 12 плоско-диагональных направлений
  Единорог (U) — 8 пространственных диагоналей
  Ферзь    (Q) — все 26 направлений (6 + 12 + 8)
  Король   (K) — те же 26 направлений, но ровно 1 шаг
  Конь     (N) — 24 L-образных прыжка (перестановки ±2,±1,0)
  Пешка    (P) — вперёд по оси Z, взятия по диагоналям

Принцип однопроходности:
  • Внешний цикл — по фигурам на доске (один раз каждая).
  • Внутренний цикл — по направлениям/смещениям (конечный набор).
  • Для скользящих фигур луч идёт строго вперёд по вектору,
    при первом блокере — break, без отката.
  • Для шаговых фигур — только проверка одного целевого поля
    на каждое смещение.
"""

from __future__ import annotations
from enum import Enum, auto
from typing import Dict, List, Tuple, Optional
from itertools import permutations


# ──────────────────────────────────────────────
#  Типы данных
# ──────────────────────────────────────────────

Pos = Tuple[int, int, int]          # (x, y, z)


class Color(Enum):
    WHITE = auto()
    BLACK = auto()

    @property
    def opponent(self) -> Color:
        return Color.BLACK if self is Color.WHITE else Color.WHITE


class PieceKind(Enum):
    KING   = "K"
    QUEEN  = "Q"
    ROOK   = "R"
    BISHOP = "B"
    UNICORN = "U"   # пространственная диагональ — классика Raumschach
    KNIGHT = "N"
    PAWN   = "P"


class Piece:
    __slots__ = ("kind", "color")

    def __init__(self, kind: PieceKind, color: Color):
        self.kind = kind
        self.color = color

    def __repr__(self) -> str:
        c = "w" if self.color is Color.WHITE else "b"
        return f"{c}{self.kind.value}"


# ──────────────────────────────────────────────
#  Предвычисленные таблицы направлений
#  (создаются один раз, переиспользуются)
# ──────────────────────────────────────────────

def _build_directions():
    """
    Генерирует все 26 единичных 3D-направлений и классифицирует их:
      осевые           (6)  — ровно одна ненулевая координата
      плоско-диагональные (12) — ровно две ненулевые координаты
      пространственные   (8)  — все три координаты ненулевые
    """
    axial = []       # Ладья
    planar = []      # Слон
    spatial = []     # Единорог

    for dx in (-1, 0, 1):
        for dy in (-1, 0, 1):
            for dz in (-1, 0, 1):
                if dx == dy == dz == 0:
                    continue
                nonzero = (dx != 0) + (dy != 0) + (dz != 0)
                d = (dx, dy, dz)
                if nonzero == 1:
                    axial.append(d)
                elif nonzero == 2:
                    planar.append(d)
                else:
                    spatial.append(d)

    return axial, planar, spatial


AXIAL_DIRS, PLANAR_DIRS, SPATIAL_DIRS = _build_directions()

ROOK_DIRS    = AXIAL_DIRS                                # 6
BISHOP_DIRS  = PLANAR_DIRS                               # 12
UNICORN_DIRS = SPATIAL_DIRS                               # 8
QUEEN_DIRS   = AXIAL_DIRS + PLANAR_DIRS + SPATIAL_DIRS    # 26
KING_DIRS    = QUEEN_DIRS                                 # 26 (шаг = 1)


def _build_knight_offsets() -> List[Pos]:
    """
    Конь в 3D: все перестановки (±2, ±1, 0) с учётом знаков.
    Множество устраняет дубли (знак на нулевой компоненте не меняет вектор).
    Результат — 24 уникальных смещения.
    """
    offsets: set[Pos] = set()
    for perm in permutations((2, 1, 0)):
        for sx in (1, -1):
            for sy in (1, -1):
                for sz in (1, -1):
                    offsets.add((perm[0] * sx, perm[1] * sy, perm[2] * sz))
    return list(offsets)


KNIGHT_OFFSETS = _build_knight_offsets()   # 24 прыжка


# ──────────────────────────────────────────────
#  Доска и генератор ходов
# ──────────────────────────────────────────────

class Board3D:
    """Трёхмерная шахматная доска N×N×N."""

    def __init__(self, size: int = 8):
        self.size = size
        self.cells: Dict[Pos, Piece] = {}

    # ── работа с доской ──

    def place(self, pos: Pos, piece: Piece) -> None:
        self.cells[pos] = piece

    def remove(self, pos: Pos) -> Optional[Piece]:
        return self.cells.pop(pos, None)

    def in_bounds(self, pos: Pos) -> bool:
        x, y, z = pos
        s = self.size
        return 0 <= x < s and 0 <= y < s and 0 <= z < s

    # ── главный метод — однопроходной сбор ходов ──

    def compute_all_moves(self) -> Dict[Pos, List[Pos]]:
        """
        Один проход по всем фигурам на доске.
        Возвращает {позиция_фигуры: [список_допустимых_полей]}.

        Сложность: O(P · D · L)
          P — число фигур на доске
          D — макс. число направлений у фигуры (≤ 26)
          L — макс. длина луча (≤ N-1)
        Для шаговых фигур L = 1, для коня D = 24, L = 1.
        """
        result: Dict[Pos, List[Pos]] = {}

        # ── единственный проход ──
        for pos, piece in self.cells.items():
            result[pos] = self._moves_for(pos, piece)

        return result

    # ── диспетчер по типу фигуры ──

    def _moves_for(self, pos: Pos, piece: Piece) -> List[Pos]:
        kind = piece.kind

        if kind is PieceKind.ROOK:
            return self._ray_moves(pos, piece, ROOK_DIRS)

        if kind is PieceKind.BISHOP:
            return self._ray_moves(pos, piece, BISHOP_DIRS)

        if kind is PieceKind.UNICORN:
            return self._ray_moves(pos, piece, UNICORN_DIRS)

        if kind is PieceKind.QUEEN:
            return self._ray_moves(pos, piece, QUEEN_DIRS)

        if kind is PieceKind.KING:
            return self._jump_moves(pos, piece, KING_DIRS)

        if kind is PieceKind.KNIGHT:
            return self._jump_moves(pos, piece, KNIGHT_OFFSETS)

        if kind is PieceKind.PAWN:
            return self._pawn_moves(pos, piece)

        return []

    # ── скользящие фигуры (луч вперёд, break при блокере) ──

    def _ray_moves(self, origin: Pos, piece: Piece,
                   directions: List[Pos]) -> List[Pos]:
        """
        Для каждого направления двигаемся строго вперёд.
        Остановка при выходе за границу или встрече фигуры.
        Фигуру противника можно взять (добавляем поле),
        свою — нельзя (не добавляем).  Оба случая — break.
        """
        moves: List[Pos] = []
        cells = self.cells
        color = piece.color
        size = self.size

        for dx, dy, dz in directions:
            x, y, z = origin

            # ── луч вдоль одного направления ──
            while True:
                x += dx
                y += dy
                z += dz

                # выход за границу → конец луча
                if not (0 <= x < size and 0 <= y < size and 0 <= z < size):
                    break

                blocker = cells.get((x, y, z))

                if blocker is None:
                    # пустое поле — можно идти, луч продолжается
                    moves.append((x, y, z))
                    continue

                if blocker.color is not color:
                    # фигура противника — взятие, луч кончается
                    moves.append((x, y, z))

                # своя фигура или только что взятая чужая — break
                break

        return moves

    # ── шаговые фигуры (один прыжок по каждому смещению) ──

    def _jump_moves(self, origin: Pos, piece: Piece,
                    offsets: List[Pos]) -> List[Pos]:
        """
        Проверяем каждое смещение ровно один раз.
        Если поле в пределах доски и не занято своей фигурой — ход возможен.
        """
        moves: List[Pos] = []
        cells = self.cells
        color = piece.color
        size = self.size
        ox, oy, oz = origin

        for dx, dy, dz in offsets:
            nx, ny, nz = ox + dx, oy + dy, oz + dz

            if not (0 <= nx < size and 0 <= ny < size and 0 <= nz < size):
                continue

            blocker = cells.get((nx, ny, nz))
            if blocker is None or blocker.color is not color:
                moves.append((nx, ny, nz))

        return moves

    # ── пешка ──

    def _pawn_moves(self, origin: Pos, piece: Piece) -> List[Pos]:
        """
        Пешка идёт вперёд по оси Z (белые +z, чёрные -z).
        Со стартового уровня — может шагнуть на 2.
        Взятия — по диагоналям в плоскости (x,y) со сдвигом +z/−z.
        """
        moves: List[Pos] = []
        cells = self.cells
        size = self.size
        ox, oy, oz = origin

        fwd = 1 if piece.color is Color.WHITE else -1
        start_z = 1 if piece.color is Color.WHITE else size - 2

        # ── тихий ход на 1 клетку вперёд ──
        z1 = oz + fwd
        if 0 <= z1 < size and (ox, oy, z1) not in cells:
            moves.append((ox, oy, z1))

            # ── двойной ход со старта ──
            if oz == start_z:
                z2 = oz + 2 * fwd
                if 0 <= z2 < size and (ox, oy, z2) not in cells:
                    moves.append((ox, oy, z2))

        # ── взятия: сдвиг по x или y (±1) + шаг по z ──
        #    В 3D пешка может бить в 8 диагональных направлениях
        #    с продвижением по z.
        capture_dxdy = [
            (-1, 0), (1, 0), (0, -1), (0, 1),    # осевые в плоскости
            (-1, -1), (-1, 1), (1, -1), (1, 1),   # диагональные в плоскости
        ]
        for dx, dy in capture_dxdy:
            cx, cy, cz = ox + dx, oy + dy, oz + fwd
            if 0 <= cx < size and 0 <= cy < size and 0 <= cz < size:
                target = cells.get((cx, cy, cz))
                if target is not None and target.color is not piece.color:
                    moves.append((cx, cy, cz))

        return moves

    # ── красивый вывод ──

    def print_moves(self, moves: Dict[Pos, List[Pos]]) -> None:
        for pos in sorted(moves):
            piece = self.cells[pos]
            targets = moves[pos]
            coords = ", ".join(str(t) for t in sorted(targets))
            print(f"  {piece} @ {pos}  →  [{len(targets)} ходов] {coords}")


# ──────────────────────────────────────────────
#  Демонстрация
# ──────────────────────────────────────────────

def demo():
    board = Board3D(size=8)

    # Расставим несколько фигур для примера
    board.place((3, 3, 3), Piece(PieceKind.QUEEN,   Color.WHITE))
    board.place((3, 3, 6), Piece(PieceKind.ROOK,    Color.BLACK))
    board.place((5, 5, 5), Piece(PieceKind.UNICORN, Color.WHITE))
    board.place((0, 0, 0), Piece(PieceKind.KNIGHT,  Color.WHITE))
    board.place((4, 4, 1), Piece(PieceKind.PAWN,    Color.WHITE))
    board.place((5, 5, 2), Piece(PieceKind.PAWN,    Color.BLACK))
    board.place((4, 4, 4), Piece(PieceKind.KING,    Color.BLACK))

    print("=" * 60)
    print("  Трёхмерные шахматы — однопроходной генератор ходов")
    print("=" * 60)
    print(f"  Доска: {board.size}×{board.size}×{board.size}")
    print(f"  Фигур на доске: {len(board.cells)}\n")

    # ── один вызов — все ходы ──
    all_moves = board.compute_all_moves()

    board.print_moves(all_moves)

    total = sum(len(m) for m in all_moves.values())
    print(f"\n  Всего допустимых ходов: {total}")
    print("=" * 60)


if __name__ == "__main__":
    demo()