/**

Copyright (c) Scott Gasch

Module Name:

    nnuefeatures.c

Abstract:

    Standalone reference implementation of the HalfKAv2_hm feature
    index formula, used ONLY to cross-check against the Python
    reference encoder (annotate_pgns/nnue_feature_spec.py) -- see
    docs/NNUE.md's "Canonical HalfKAv2_hm feature-indexing spec"
    section for the formulas this file implements, sourced verbatim
    from Stockfish's src/nnue/features/half_ka_v2_hm.{h,cpp} (git tag
    sf_16).

    This is deliberately NOT wired into search/eval -- no accumulator,
    no incremental maintenance, nothing performance-sensitive. It
    exists purely so `dump nnuefeatures` can print the exact same
    active-feature-index set the eventual real C inference code must
    produce, so a Python/C cross-checker has a C-side ground truth to
    diff against before any of that real code exists. The *real*
    MakeMove/UnmakeMove accumulator plumbing (per NNUE.md's
    sequencing, gated behind a dedicated branch, not started) will
    reuse this same index math, not this file's dump-command
    plumbing.

Author:

    Scott Gasch (scott.gasch@gmail.com) 15 Oct 2026

**/

#include "chess.h"

//
// HalfKAv2_hm constants, verbatim from Stockfish's half_ka_v2_hm.h.
// PS_NB counts 5 non-king piece types x 2 "us/them" roles x 64
// squares, plus one shared 64-square king plane (own king and enemy
// king features both land in this one plane, distinguished only by
// which square is set).
//
#define NNUE_SQUARE_NB    (64)
#define NNUE_PS_W_PAWN    (0 * NNUE_SQUARE_NB)
#define NNUE_PS_B_PAWN    (1 * NNUE_SQUARE_NB)
#define NNUE_PS_W_KNIGHT  (2 * NNUE_SQUARE_NB)
#define NNUE_PS_B_KNIGHT  (3 * NNUE_SQUARE_NB)
#define NNUE_PS_W_BISHOP  (4 * NNUE_SQUARE_NB)
#define NNUE_PS_B_BISHOP  (5 * NNUE_SQUARE_NB)
#define NNUE_PS_W_ROOK    (6 * NNUE_SQUARE_NB)
#define NNUE_PS_B_ROOK    (7 * NNUE_SQUARE_NB)
#define NNUE_PS_W_QUEEN   (8 * NNUE_SQUARE_NB)
#define NNUE_PS_B_QUEEN   (9 * NNUE_SQUARE_NB)
#define NNUE_PS_KING      (10 * NNUE_SQUARE_NB)
#define NNUE_PS_NB        (11 * NNUE_SQUARE_NB)        // 704

#define NNUE_DIMENSIONS   ((NNUE_SQUARE_NB * NNUE_PS_NB) / 2)   // 22528

//
// Convert a typhoon 0x88 COOR to Stockfish's flat 0-63 square index
// (file + 8*(rank-1), rank/file both derived from the engine's own
// RANK()/FILE() macros -- see docs/NNUE.md's spec for why this
// one-line conversion is the single highest-risk spot in the whole
// scheme).
//
static int
CoorToSq64(COOR c)
{
    return (int)FILE(c) + 8 * ((int)RANK(c) - 1);
}

//
// King bucket (0..31) for a king on Stockfish-square sq64, from
// color uPerspective's (WHITE/BLACK) point of view. Verified against
// Stockfish's literal KingBuckets[][] table -- see docs/NNUE.md.
//
static int
KingBucket(ULONG uPerspective, int sq64)
{
    int file = sq64 & 7;
    int rank0 = sq64 >> 3;              // 0..7, rank-1
    int fileQuartet = (file < 4) ? file : (7 - file);   // 0..3, mirrored
    int rankGroup = (uPerspective == WHITE) ? (7 - rank0) : rank0;
    return rankGroup * 4 + fileQuartet;
}

//
// Orientation XOR mask for perspective uPerspective whose own king
// sits on Stockfish-square ksq64. Only ksq64's file selects the
// branch (queenside half a-d vs. kingside half e-h); any rank.
//
static int
OrientMask(ULONG uPerspective, int ksq64)
{
    int file = ksq64 & 7;
    if (uPerspective == WHITE)
    {
        return (file < 4) ? 7 : 0;          // SQ_H1 : SQ_A1
    }
    return (file < 4) ? 63 : 56;            // SQ_H8 : SQ_A8
}

//
// Piece-square plane offset for a piece of color uPieceColor and
// type uPieceType (chess.h's PAWN..KING values), as seen from
// perspective uPerspective. Both kings (own and enemy) share
// NNUE_PS_KING.
//
static int
PiecePlane(ULONG uPerspective, ULONG uPieceColor, ULONG uPieceType)
{
    FLAG fOwn;

    if (uPieceType == KING)
    {
        return NNUE_PS_KING;
    }
    fOwn = (uPieceColor == uPerspective);
    switch (uPieceType)
    {
        case PAWN:   return fOwn ? NNUE_PS_W_PAWN   : NNUE_PS_B_PAWN;
        case KNIGHT: return fOwn ? NNUE_PS_W_KNIGHT : NNUE_PS_B_KNIGHT;
        case BISHOP: return fOwn ? NNUE_PS_W_BISHOP : NNUE_PS_B_BISHOP;
        case ROOK:   return fOwn ? NNUE_PS_W_ROOK   : NNUE_PS_B_ROOK;
        case QUEEN:  return fOwn ? NNUE_PS_W_QUEEN  : NNUE_PS_B_QUEEN;
        default:
            ASSERT(FALSE);
            return 0;
    }
}

static int
CompareInts(const void *a, const void *b)
{
    int x = *(const int *)a;
    int y = *(const int *)b;
    return (x > y) - (x < y);
}

//
// Fill pIndices (caller-allocated, at least 32 ints) with the sorted
// set of active HalfKAv2_hm feature indices for perspective
// uPerspective at position pos. Returns the count (<=32).
//
static int
ActiveFeatureIndices(POSITION *pos, ULONG uPerspective, int *pIndices)
{
    int f, r;
    COOR cKing = ILLEGAL_COOR;
    int ksq64;
    int orientMask;
    int kingBucketOffset;
    int n = 0;

    //
    // Find this perspective's own king square.
    //
    for (f = 0; f < 8; f++)
    {
        for (r = 1; r <= 8; r++)
        {
            COOR c = FILE_RANK_TO_COOR(f, r);
            PIECE p = pos->rgSquare[c].pPiece;
            if (!IS_EMPTY(p) &&
                (PIECE_TYPE(p) == KING) &&
                (PIECE_COLOR(p) == uPerspective))
            {
                cKing = c;
            }
        }
    }
    ASSERT(cKing != ILLEGAL_COOR);
    ksq64 = CoorToSq64(cKing);
    orientMask = OrientMask(uPerspective, ksq64);
    kingBucketOffset = KingBucket(uPerspective, ksq64) * NNUE_PS_NB;

    for (f = 0; f < 8; f++)
    {
        for (r = 1; r <= 8; r++)
        {
            COOR c = FILE_RANK_TO_COOR(f, r);
            PIECE p = pos->rgSquare[c].pPiece;
            int sq64;
            int idx;

            if (IS_EMPTY(p))
            {
                continue;
            }
            sq64 = CoorToSq64(c);
            idx = (sq64 ^ orientMask) +
                  PiecePlane(uPerspective, PIECE_COLOR(p), PIECE_TYPE(p)) +
                  kingBucketOffset;
            ASSERT((idx >= 0) && (idx < NNUE_DIMENSIONS));
            ASSERT(n < 32);
            pIndices[n++] = idx;
        }
    }
    qsort(pIndices, n, sizeof(int), CompareInts);
    return n;
}

void
DumpNnueFeatures(POSITION *pos)
/**

Routine description:

    Prints the sorted set of active HalfKAv2_hm feature indices for
    both perspectives at the current position, one line per
    perspective, in a format trivial for a Python cross-checker to
    parse:

        NNUE_FEATURES white <idx> <idx> ...
        NNUE_FEATURES black <idx> <idx> ...

Parameters:

    POSITION *pos : position to dump

Return value:

    void

**/
{
    int rgIndices[32];
    int n;
    int i;

    n = ActiveFeatureIndices(pos, WHITE, rgIndices);
    Trace("NNUE_FEATURES white");
    for (i = 0; i < n; i++)
    {
        Trace(" %d", rgIndices[i]);
    }
    Trace("\n");

    n = ActiveFeatureIndices(pos, BLACK, rgIndices);
    Trace("NNUE_FEATURES black");
    for (i = 0; i < n; i++)
    {
        Trace(" %d", rgIndices[i]);
    }
    Trace("\n");
}