from __future__ import annotations

from dataclasses import dataclass


@dataclass(frozen=True)
class Piece:
    """A polyomino piece defined by a set of (x, y) coordinates.

    Coordinates are relative to a reference point (typically the
    top-left corner of the bounding box after normalization).
    """

    name: str
    squares: frozenset[tuple[int, int]]

    @property
    def size(self) -> int:
        return len(self.squares)

    def to_dict(self) -> dict:
        return {
            "name": self.name,
            "squares": sorted(self.squares),
        }

    @classmethod
    def from_dict(cls, data: dict) -> Piece:
        return cls(
            name=data["name"],
            squares=frozenset(tuple(s) for s in data["squares"]),
        )


@dataclass
class OrientedPiece:
    """A piece in a specific orientation placed at a specific position."""

    piece_id: int
    orientation_id: int
    squares: list[tuple[int, int]]
    corners: list[tuple[int, int]]


@dataclass
class Move:
    """A complete move: which piece, which orientation, where placed."""

    piece_id: int
    orientation_id: int
    x: int
    y: int


@dataclass
class PieceOrientation:
    """A piece in a specific orientation (before placement on board)."""

    piece_id: int
    orientation_id: int
    squares: list[tuple[int, int]]  # normalized relative coordinates
    corners: list[tuple[int, int]]  # corner cells relative to squares


def _normalize(squares: list[tuple[int, int]]) -> list[tuple[int, int]]:
    """Shift coordinates so the minimum x and y are 0."""
    min_x = min(x for x, _ in squares)
    min_y = min(y for _, y in squares)
    return [(x - min_x, y - min_y) for x, y in squares]


def _rotate(squares: list[tuple[int, int]]) -> list[tuple[int, int]]:
    """Rotate 90 degrees clockwise: (x, y) -> (y, -x)."""
    return [(y, -x) for x, y in squares]


def _flip(squares: list[tuple[int, int]]) -> list[tuple[int, int]]:
    """Flip horizontally: (x, y) -> (-x, y)."""
    return [(-x, y) for x, y in squares]


def _compute_corners(squares: list[tuple[int, int]]) -> list[tuple[int, int]]:
    """Compute the corner cells for a piece.

    Corner cells are the diagonal neighbors of each square that are not
    edge-adjacent to any other square in the piece. These are the cells
    where a same-color piece must touch (corner-to-corner).
    """
    square_set = set(squares)
    corners = set()
    for x, y in squares:
        for dx, dy in [(-1, -1), (1, -1), (-1, 1), (1, 1)]:
            cx, cy = x + dx, y + dy
            if (cx, cy) not in square_set:
                # Check that this corner is not edge-adjacent to another square
                edge_neighbors = [
                    (cx + 1, cy), (cx - 1, cy), (cx, cy + 1), (cx, cy - 1)
                ]
                if not any(en in square_set for en in edge_neighbors):
                    corners.add((cx, cy))
    return sorted(corners)


def generate_orientations(piece: Piece) -> list[PieceOrientation]:
    """Generate all unique orientations of a piece.

    Applies 4 rotations and 2 flips, then deduplicates by comparing
    normalized coordinate sets.
    """
    orientations: list[PieceOrientation] = []
    seen: set[tuple[tuple[int, int], ...]] = set()

    squares = sorted(piece.squares)
    for flip_count in range(2):
        for rotation_count in range(4):
            current = list(squares)
            for _ in range(rotation_count):
                current = _rotate(current)
            for _ in range(flip_count):
                current = _flip(current)
            normalized = tuple(sorted(_normalize(current)))
            if normalized not in seen:
                seen.add(normalized)
                corners = _compute_corners(list(normalized))
                orientations.append(
                    PieceOrientation(
                        piece_id=-1,  # Set by PieceSet
                        orientation_id=len(orientations),
                        squares=list(normalized),
                        corners=corners,
                    )
                )
    return orientations


# ---------------------------------------------------------------------------
# Standard Blokus piece sets
# ---------------------------------------------------------------------------

# The 21 free polyominoes of size 1-5
# Naming follows standard pentomino/tetromino conventions

_I1 = Piece("I1", frozenset([(0, 0)]))

_I2 = Piece("I2", frozenset([(0, 0), (0, 1)]))

_I3 = Piece("I3", frozenset([(0, 0), (0, 1), (0, 2)]))

_V3 = Piece("V3", frozenset([(0, 0), (1, 0), (0, 1)]))

_I4 = Piece("I4", frozenset([(0, 0), (0, 1), (0, 2), (0, 3)]))

_L4 = Piece("L4", frozenset([(0, 0), (0, 1), (0, 2), (1, 0)]))

_T4 = Piece("T4", frozenset([(0, 0), (1, 0), (2, 0), (1, 1)]))

_S4 = Piece("S4", frozenset([(0, 0), (1, 0), (1, 1), (2, 1)]))

_O4 = Piece("O4", frozenset([(0, 0), (1, 0), (0, 1), (1, 1)]))

_I5 = Piece("I5", frozenset([(0, 0), (0, 1), (0, 2), (0, 3), (0, 4)]))

_L5 = Piece("L5", frozenset([(0, 0), (0, 1), (0, 2), (0, 3), (1, 0)]))

_Y5 = Piece("Y5", frozenset([(0, 0), (0, 1), (0, 2), (0, 3), (1, 1)]))

_N5 = Piece("N5", frozenset([(0, 0), (1, 0), (2, 0), (2, 1), (3, 1)]))

_T5 = Piece("T5", frozenset([(0, 0), (1, 0), (2, 0), (1, 1), (1, 2)]))

_U5 = Piece("U5", frozenset([(0, 0), (2, 0), (0, 1), (1, 1), (2, 1)]))

_V5 = Piece("V5", frozenset([(0, 0), (0, 1), (0, 2), (1, 0), (2, 0)]))

_W5 = Piece("W5", frozenset([(0, 0), (0, 1), (1, 0), (1, 1), (2, 0)]))

_Z5 = Piece("Z5", frozenset([(0, 0), (1, 0), (1, 1), (1, 2), (2, 2)]))

_F5 = Piece("F5", frozenset([(0, 0), (1, 0), (1, 1), (2, 1), (1, 2)]))

_X5 = Piece("X5", frozenset([(0, 0), (1, 0), (2, 0), (1, 1), (1, -1)]))

_P5 = Piece("P5", frozenset([(0, 0), (1, 0), (0, 1), (1, 1), (0, 2)]))

STANDARD_PIECES: list[Piece] = [
    _I1, _I2, _I3, _V3, _I4, _L4, _T4, _S4, _O4,
    _I5, _L5, _Y5, _N5, _T5, _U5, _V5, _W5, _Z5, _F5, _X5, _P5,
]

# Blokus Duo uses the same pieces but on a 14x14 board
DUO_PIECES: list[Piece] = STANDARD_PIECES

# Blokus Junior: simplified pieces (only 12 unique pieces, 2 copies each)
# Uses only pieces with size <= 4 for simplicity
JUNIOR_PIECES: list[Piece] = [
    _I1, _I2, _I3, _V3, _I4, _L4, _T4, _S4, _O4,
]


class PieceSet:
    """A collection of pieces with pre-computed orientations.

    This class manages the piece set used in a Blokus game, including
    all unique orientations for each piece and a lookup table for
    generating moves.
    """

    def __init__(self, pieces: list[Piece]):
        self.pieces: list[Piece] = pieces
        self.piece_names: list[str] = [p.name for p in pieces]
        self.piece_id_map: dict[str, int] = {name: i for i, name in enumerate(self.piece_names)}

        # Pre-compute all orientations for each piece
        self.orientations: list[list[PieceOrientation]] = []
        for piece_id, piece in enumerate(pieces):
            orients = generate_orientations(piece)
            for orient in orients:
                orient.piece_id = piece_id
            self.orientations.append(orients)

        # Total number of (piece, orientation) combinations
        self.num_orientations: int = sum(len(o) for o in self.orientations)

    @property
    def num_pieces(self) -> int:
        return len(self.pieces)

    def get_orientations(self, piece_id: int) -> list[PieceOrientation]:
        return self.orientations[piece_id]

    def get_piece(self, piece_id: int) -> Piece:
        return self.pieces[piece_id]

    def get_piece_id(self, name: str) -> int:
        return self.piece_id_map[name]

    def to_dict(self) -> dict:
        return {
            "pieces": [p.to_dict() for p in self.pieces],
        }

    @classmethod
    def from_dict(cls, data: dict) -> PieceSet:
        pieces = [Piece.from_dict(p) for p in data["pieces"]]
        return cls(pieces)

    @classmethod
    def from_json(cls, path: str) -> PieceSet:
        import json

        with open(path) as f:
            data = json.load(f)
        return cls.from_dict(data)

    def save_json(self, path: str) -> None:
        import json

        with open(path, "w") as f:
            json.dump(self.to_dict(), f, indent=2)

    def __len__(self) -> int:
        return len(self.pieces)

    def __repr__(self) -> str:
        return f"PieceSet(num_pieces={self.num_pieces}, num_orientations={self.num_orientations})"