"""
Board and move encoding for TinyChess.

Design notes
------------
* Positions are *canonicalised* to side-to-move = White by vertically mirroring
  the board and swapping piece colours when Black is to move. This is a
  lossless symmetry of chess (ignoring nothing: castling rights and en-passant
  are mirrored too), and it is a standard, non-cheating canonicalisation.
* A move is packed into a single int32 with structural fields only. We
  deliberately store *rules-derived* structure (from/to/promo/piece/capture/
  castle/ep) and no heuristic evaluation.
* Per-square features are kept separate (64 vectors); nothing is collapsed to a
  single vector at encoding time.
"""
from __future__ import annotations

import numpy as np
import chess

# ---------------------------------------------------------------------------
# Move packing
# ---------------------------------------------------------------------------
# bits: from(6) | to(6) | promo(3) | piece(3) | capture(1) | castle(1) | ep(1)
_F_FROM = 0
_F_TO = 6
_F_PROMO = 12
_F_PIECE = 15
_F_CAP = 18
_F_CASTLE = 19
_F_EP = 20

PROMO_CODES = {None: 0, chess.KNIGHT: 1, chess.BISHOP: 2, chess.ROOK: 3, chess.QUEEN: 4}
PROMO_INV = {v: k for k, v in PROMO_CODES.items()}
N_PROMO = 5
N_PIECE_TYPES = 6  # pawn..king -> 0..5

MOVE_FIELD_SIZES = dict(from_sq=64, to_sq=64, promo=N_PROMO, piece=N_PIECE_TYPES,
                        capture=2, castle=2, ep=2)


def pack_move(board: chess.Board, move: chess.Move) -> int:
    """Pack a legal move on `board` into an int32 of structural fields."""
    piece = board.piece_at(move.from_square)
    ptype = (piece.piece_type - 1) if piece is not None else 0
    is_ep = board.is_en_passant(move)
    is_cap = board.is_capture(move)
    is_castle = board.is_castling(move)
    v = (move.from_square << _F_FROM) | (move.to_square << _F_TO)
    v |= PROMO_CODES.get(move.promotion, 0) << _F_PROMO
    v |= ptype << _F_PIECE
    v |= int(is_cap) << _F_CAP
    v |= int(is_castle) << _F_CASTLE
    v |= int(is_ep) << _F_EP
    return v


def unpack_move(v: int) -> dict:
    return dict(
        from_sq=(v >> _F_FROM) & 63,
        to_sq=(v >> _F_TO) & 63,
        promo=(v >> _F_PROMO) & 7,
        piece=(v >> _F_PIECE) & 7,
        capture=(v >> _F_CAP) & 1,
        castle=(v >> _F_CASTLE) & 1,
        ep=(v >> _F_EP) & 1,
    )


def unpack_fields(v: np.ndarray) -> np.ndarray:
    """Vectorised unpack -> int64 array [..., 7] in MOVE_FIELD order."""
    v = np.asarray(v, dtype=np.int64)
    return np.stack([
        (v >> _F_FROM) & 63,
        (v >> _F_TO) & 63,
        (v >> _F_PROMO) & 7,
        (v >> _F_PIECE) & 7,
        (v >> _F_CAP) & 1,
        (v >> _F_CASTLE) & 1,
        (v >> _F_EP) & 1,
    ], axis=-1)


MOVE_FIELD_ORDER = ["from_sq", "to_sq", "promo", "piece", "capture", "castle", "ep"]


def packed_to_uci(v: int) -> str:
    d = unpack_move(v)
    m = chess.Move(d["from_sq"], d["to_sq"], promotion=PROMO_INV.get(d["promo"]))
    return m.uci()


def packed_to_move(v: int) -> chess.Move:
    d = unpack_move(v)
    return chess.Move(d["from_sq"], d["to_sq"], promotion=PROMO_INV.get(d["promo"]))


# ---------------------------------------------------------------------------
# Canonicalisation
# ---------------------------------------------------------------------------

def canonical_board(board: chess.Board) -> tuple[chess.Board, bool]:
    """Return (board with White to move, flipped?).

    Uses python-chess `mirror()` which flips vertically and swaps colours,
    correctly transforming castling rights and the en-passant square.
    """
    if board.turn == chess.WHITE:
        return board, False
    return board.mirror(), True


def flip_square(sq: int) -> int:
    return sq ^ 56


def flip_move(move: chess.Move) -> chess.Move:
    return chess.Move(flip_square(move.from_square), flip_square(move.to_square),
                      promotion=move.promotion)


# ---------------------------------------------------------------------------
# Board -> feature arrays
# ---------------------------------------------------------------------------
# Per-square categorical: 0 = empty, 1..6 = own P,N,B,R,Q,K, 7..12 = opp
N_PIECE_TOKENS = 13
# Per-square binary extras
SQ_EXTRA = ["last_from", "last_to", "ep_square", "own_backrank_castle"]
N_SQ_EXTRA = len(SQ_EXTRA)

GLOBAL_FEATS = [
    "own_oo", "own_ooo", "opp_oo", "opp_ooo",
    "has_ep",
    "halfmove_clock_n", "fullmove_n", "repetition_n",
    "own_pawns_n", "own_knights_n", "own_bishops_n", "own_rooks_n", "own_queens_n",
    "opp_pawns_n", "opp_knights_n", "opp_bishops_n", "opp_rooks_n", "opp_queens_n",
    "in_check", "material_diff_n",
]
N_GLOBAL = len(GLOBAL_FEATS)

_PT_ORDER = [chess.PAWN, chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN, chess.KING]
_PIECE_VAL = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5,
              chess.QUEEN: 9, chess.KING: 0}


def encode_board(cboard: chess.Board, last_move: chess.Move | None = None,
                 repetition: int = 0):
    """Encode a *canonical* (White-to-move) board.

    Returns (squares int8[64], extras int8[64,N_SQ_EXTRA], glob float32[N_GLOBAL]).
    """
    sq = np.zeros(64, dtype=np.int8)
    for s, piece in cboard.piece_map().items():
        idx = _PT_ORDER.index(piece.piece_type) + 1
        if piece.color == chess.BLACK:
            idx += 6
        sq[s] = idx

    extras = np.zeros((64, N_SQ_EXTRA), dtype=np.int8)
    if last_move is not None:
        extras[last_move.from_square, 0] = 1
        extras[last_move.to_square, 1] = 1
    if cboard.ep_square is not None:
        extras[cboard.ep_square, 2] = 1
    if cboard.has_kingside_castling_rights(chess.WHITE):
        extras[chess.H1, 3] = 1
        extras[chess.E1, 3] = 1
    if cboard.has_queenside_castling_rights(chess.WHITE):
        extras[chess.A1, 3] = 1
        extras[chess.E1, 3] = 1

    g = np.zeros(N_GLOBAL, dtype=np.float32)
    g[0] = cboard.has_kingside_castling_rights(chess.WHITE)
    g[1] = cboard.has_queenside_castling_rights(chess.WHITE)
    g[2] = cboard.has_kingside_castling_rights(chess.BLACK)
    g[3] = cboard.has_queenside_castling_rights(chess.BLACK)
    g[4] = cboard.ep_square is not None
    g[5] = min(cboard.halfmove_clock, 100) / 100.0
    g[6] = min(cboard.fullmove_number, 120) / 120.0
    g[7] = min(repetition, 3) / 3.0
    mat = 0
    for i, pt in enumerate(_PT_ORDER[:5]):
        nw = len(cboard.pieces(pt, chess.WHITE))
        nb = len(cboard.pieces(pt, chess.BLACK))
        g[8 + i] = nw / 8.0
        g[13 + i] = nb / 8.0
        mat += _PIECE_VAL[pt] * (nw - nb)
    g[18] = cboard.is_check()
    g[19] = np.tanh(mat / 10.0)
    return sq, extras, g


def encode_position(board: chess.Board, last_move: chess.Move | None = None,
                    repetition: int = 0):
    """Canonicalise then encode. Returns (squares, extras, glob, flipped)."""
    cb, flipped = canonical_board(board)
    lm = last_move
    if flipped and lm is not None:
        lm = flip_move(lm)
    sq, ex, g = encode_board(cb, lm, repetition)
    return sq, ex, g, flipped


def legal_packed(board: chess.Board):
    """Return (packed candidate array int32[L], canonical board, flipped)."""
    cb, flipped = canonical_board(board)
    return np.array([pack_move(cb, m) for m in cb.legal_moves], dtype=np.int32), cb, flipped