#!/usr/bin/env python3
"""Python reference implementation of the HalfKAv2_hm feature index
formula -- see ../src/docs/NNUE.md's "Canonical HalfKAv2_hm
feature-indexing spec" section, sourced verbatim from Stockfish's
src/nnue/features/half_ka_v2_hm.{h,cpp} (git tag sf_16).

This module's only job is to be an independent, from-scratch
implementation of the same math as ../src/nnuefeatures.c
(`dump nnuefeatures`), using python-chess's own (already
rank*8+file, 0-63, a1=0) square numbering directly -- no 0x88
conversion needed on this side, that's purely a typhoon-C-side
concern.

See cross_check_features.py for the actual Python-vs-C diff driver.
"""
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src" / "eval_tune"))
import chess  # noqa: E402

SQUARE_NB = 64

PS_W_PAWN, PS_B_PAWN = 0 * SQUARE_NB, 1 * SQUARE_NB
PS_W_KNIGHT, PS_B_KNIGHT = 2 * SQUARE_NB, 3 * SQUARE_NB
PS_W_BISHOP, PS_B_BISHOP = 4 * SQUARE_NB, 5 * SQUARE_NB
PS_W_ROOK, PS_B_ROOK = 6 * SQUARE_NB, 7 * SQUARE_NB
PS_W_QUEEN, PS_B_QUEEN = 8 * SQUARE_NB, 9 * SQUARE_NB
PS_KING = 10 * SQUARE_NB
PS_NB = 11 * SQUARE_NB  # 704

DIMENSIONS = (SQUARE_NB * PS_NB) // 2  # 22528

_OWN_PLANE = {
    chess.PAWN: PS_W_PAWN,
    chess.KNIGHT: PS_W_KNIGHT,
    chess.BISHOP: PS_W_BISHOP,
    chess.ROOK: PS_W_ROOK,
    chess.QUEEN: PS_W_QUEEN,
}
_ENEMY_PLANE = {
    chess.PAWN: PS_B_PAWN,
    chess.KNIGHT: PS_B_KNIGHT,
    chess.BISHOP: PS_B_BISHOP,
    chess.ROOK: PS_B_ROOK,
    chess.QUEEN: PS_B_QUEEN,
}


def king_bucket(perspective, sq):
    """sq, perspective: python-chess 0-63 square / chess.WHITE|BLACK."""
    file = chess.square_file(sq)
    rank0 = chess.square_rank(sq)  # 0..7
    file_quartet = file if file < 4 else (7 - file)
    rank_group = (7 - rank0) if perspective == chess.WHITE else rank0
    return rank_group * 4 + file_quartet


def orient_mask(perspective, ksq):
    file = chess.square_file(ksq)
    if perspective == chess.WHITE:
        return 7 if file < 4 else 0       # SQ_H1 : SQ_A1
    return 63 if file < 4 else 56          # SQ_H8 : SQ_A8


def piece_plane(perspective, piece_color, piece_type):
    if piece_type == chess.KING:
        return PS_KING
    table = _OWN_PLANE if piece_color == perspective else _ENEMY_PLANE
    return table[piece_type]


def active_feature_indices(board, perspective):
    """board: chess.Board. perspective: chess.WHITE or chess.BLACK.
    Returns a sorted list of active feature indices (len <= 32)."""
    ksq = board.king(perspective)
    mask = orient_mask(perspective, ksq)
    bucket_offset = king_bucket(perspective, ksq) * PS_NB

    indices = []
    for sq, piece in board.piece_map().items():
        idx = (sq ^ mask) + piece_plane(perspective, piece.color, piece.piece_type) + bucket_offset
        assert 0 <= idx < DIMENSIONS, idx
        indices.append(idx)
    indices.sort()
    return indices


if __name__ == "__main__":
    # Quick manual smoke test against the starting position.
    board = chess.Board()
    print("white:", active_feature_indices(board, chess.WHITE))
    print("black:", active_feature_indices(board, chess.BLACK))