#!/usr/bin/env python3
"""Chess encoding: board positions and moves as integer sequences.
No external tokenizer needed — pure chess-aware encoding.
"""
import chess

# Piece values (0 = empty, 1-6 = white, 7-12 = black)
PIECE_VALUES = {
    None: 0,
    chess.Piece(chess.PAWN, chess.WHITE): 1,
    chess.Piece(chess.KNIGHT, chess.WHITE): 2,
    chess.Piece(chess.BISHOP, chess.WHITE): 3,
    chess.Piece(chess.ROOK, chess.WHITE): 4,
    chess.Piece(chess.QUEEN, chess.WHITE): 5,
    chess.Piece(chess.KING, chess.WHITE): 6,
    chess.Piece(chess.PAWN, chess.BLACK): 7,
    chess.Piece(chess.KNIGHT, chess.BLACK): 8,
    chess.Piece(chess.BISHOP, chess.BLACK): 9,
    chess.Piece(chess.ROOK, chess.BLACK): 10,
    chess.Piece(chess.QUEEN, chess.BLACK): 11,
    chess.Piece(chess.KING, chess.BLACK): 12,
}

# Special tokens
PAD_TOKEN = 0
BOS_TOKEN = 13
EOS_TOKEN = 14
SEP_TOKEN = 15
SIDE_WHITE_TOKEN = 16
SIDE_BLACK_TOKEN = 17
MOVE_FROM_OFFSET = 18   # 18 + square (0-63) = from-square token
MOVE_TO_OFFSET = 82     # 82 + square (0-63) = to-square token
# Total vocab: 82 + 64 = 146 + promotions
# Promotion tokens: 146 + promo_type(0-3) * 64 + to_square
PROMO_OFFSET = 146       # 146 + promo*64 + to_square
# promo: 0=queen, 1=rook, 2=bishop, 3=knight
VOCAB_SIZE = PROMO_OFFSET + 4 * 64  # = 402


def encode_board(board):
    """Encode chess board to token sequence.
    Returns list of ints: [BOS, piece(64x), side, SEP]
    """
    tokens = [BOS_TOKEN]
    for square in chess.SQUARES:
        piece = board.piece_at(square)
        tokens.append(PIECE_VALUES[piece])
    tokens.append(SIDE_WHITE_TOKEN if board.turn == chess.WHITE else SIDE_BLACK_TOKEN)
    tokens.append(SEP_TOKEN)
    return tokens


def encode_move(move):
    """Encode a move to token sequence.
    Returns [from_square_token, to_square_token] or [from_square_token, promo_token]
    """
    from_sq = move.from_square
    to_sq = move.to_square
    tokens = [MOVE_FROM_OFFSET + from_sq]

    if move.promotion is not None:
        promo_map = {chess.QUEEN: 0, chess.ROOK: 1, chess.BISHOP: 2, chess.KNIGHT: 3}
        promo_type = promo_map[move.promotion]
        tokens.append(PROMO_OFFSET + promo_type * 64 + to_sq)
    else:
        tokens.append(MOVE_TO_OFFSET + to_sq)

    return tokens


def decode_move(tokens, board):
    """Decode token sequence back to a chess.Move.
    Input: list of token IDs (just the move tokens, not board tokens).
    """
    if len(tokens) < 2:
        return None

    from_token = tokens[0]
    to_token = tokens[1]

    from_sq = from_token - MOVE_FROM_OFFSET
    if from_sq < 0 or from_sq > 63:
        return None

    if to_token >= PROMO_OFFSET:
        promo_offset = to_token - PROMO_OFFSET
        promo_type = promo_offset // 64
        to_sq = promo_offset % 64
        promo_map = {0: chess.QUEEN, 1: chess.ROOK, 2: chess.BISHOP, 3: chess.KNIGHT}
        promotion = promo_map.get(promo_type)
        move = chess.Move(from_sq, to_sq, promotion=promotion)
    else:
        to_sq = to_token - MOVE_TO_OFFSET
        if to_sq < 0 or to_sq > 63:
            return None
        move = chess.Move(from_sq, to_sq)

    if move in board.legal_moves:
        return move
    return None


def encode_game(board, move):
    """Encode a board+move pair for training.
    Input sequence: board_tokens + move_tokens
    Target: move_tokens (shifted by 1)
    """
    board_tokens = encode_board(board)
    move_tokens = encode_move(move)
    full = board_tokens + move_tokens
    return full


def decode_position_to_tokens(board):
    """Get board encoding for inference."""
    return encode_board(board)