"""Hugging Face configuration for the chess transition policy."""

from __future__ import annotations

from transformers import PretrainedConfig


class ChessPolicyConfig(PretrainedConfig):
    """Serializable configuration for ``ChessTransitionPolicy``."""

    model_type = "chess_transition_policy"

    def __init__(
        self,
        *,
        board_squares: int = 64,
        piece_states: int = 13,
        d_model: int = 384,
        n_heads: int = 6,
        board_layers: int = 6,
        candidate_layers: int = 12,
        swiglu_hidden: int = 1040,
        dropout: float = 0.0,
        activation_checkpointing: bool = False,
        **kwargs,
    ) -> None:
        super().__init__(**kwargs)
        if board_squares != 64:
            raise ValueError("ChessPolicyConfig requires exactly 64 board squares")
        if d_model % n_heads != 0:
            raise ValueError("d_model must be divisible by n_heads")
        if board_layers < 1 or candidate_layers < 1:
            raise ValueError("Both transformer stacks need at least one layer")
        self.board_squares = board_squares
        self.piece_states = piece_states
        self.d_model = d_model
        self.n_heads = n_heads
        self.board_layers = board_layers
        self.candidate_layers = candidate_layers
        self.swiglu_hidden = swiglu_hidden
        self.dropout = dropout
        self.activation_checkpointing = activation_checkpointing

    @property
    def head_dim(self) -> int:
        return self.d_model // self.n_heads