/**

Copyright (c) Scott Gasch

Module Name:

    nnueeval.c

Abstract:

    The FC head (2*L1 -> L2 -> L3 -> 1) that sits on top of the two
    accumulator halves maintained by nnueaccum.c, plus the loader for
    a real trained checkpoint exported by annotate_pgns/
    export_nnue_weights.py. Reference-only right now -- not wired
    into Eval()/search, just enough to prove weights can actually be
    loaded and forward-propagated against a real position (see
    command.c's "nnue load"/"dump nnueeval").

    Not hot-path (unlike nnueaccum.c's Add/Remove, called from every
    MakeMove) -- this only runs once per position actually evaluated,
    so it's written straight-line/unoptimized, same spirit as
    nnuefeatures.c's reference dump.

Author:

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

**/

#include "chess.h"

float g_rgfNnueFc1Weight[NNUE_L2][2 * NNUE_ACCUMULATOR_SIZE];
float g_rgfNnueFc1Bias[NNUE_L2];
float g_rgfNnueFc2Weight[NNUE_L3][NNUE_L2];
float g_rgfNnueFc2Bias[NNUE_L3];
float g_rgfNnueFc3Weight[NNUE_L3];
float g_rgfNnueFc3Bias;
FLAG g_fNnueWeightsLoaded = FALSE;
FLAG g_fUseNnueEval = FALSE;


FLAG
NnueLoadWeights(CHAR *szFilename)
/**

Routine description:

    Loads a weights file written by annotate_pgns/
    export_nnue_weights.py: a 4-byte magic, a 4xuint32 dimension
    header (sanity-checked against this build's own NNUE_DIMENSIONS/
    NNUE_ACCUMULATOR_SIZE/NNUE_L2/NNUE_L3 so a mismatched export can't
    silently misread), then the eight weight/bias arrays in a fixed
    order. g_fNnueWeightsLoaded is FALSE on any failure (including a
    partial read -- no half-loaded state is left live).

Parameters:

    CHAR *szFilename

Return value:

    FLAG : TRUE on success

**/
{
    FILE *fp;
    CHAR szMagic[4];
    ULONG rguDims[4];

    g_fNnueWeightsLoaded = FALSE;
    fp = fopen(szFilename, "rb");
    if (NULL == fp)
    {
        Trace("NnueLoadWeights: could not open %s\n", szFilename);
        return FALSE;
    }

    if ((1 != fread(szMagic, 4, 1, fp)) ||
        (0 != memcmp(szMagic, "NNU1", 4)))
    {
        Trace("NnueLoadWeights: %s is not a valid NNUE weights file "
              "(bad magic)\n", szFilename);
        fclose(fp);
        return FALSE;
    }

    if (1 != fread(rguDims, sizeof(rguDims), 1, fp))
    {
        Trace("NnueLoadWeights: %s truncated (header)\n", szFilename);
        fclose(fp);
        return FALSE;
    }
    if ((rguDims[0] != NNUE_DIMENSIONS) ||
        (rguDims[1] != NNUE_ACCUMULATOR_SIZE) ||
        (rguDims[2] != NNUE_L2) ||
        (rguDims[3] != NNUE_L3))
    {
        Trace("NnueLoadWeights: %s dims (%lu,%lu,%lu,%lu) don't match "
              "this build's (%lu,%lu,%lu,%lu)\n", szFilename,
              rguDims[0], rguDims[1], rguDims[2], rguDims[3],
              (ULONG)NNUE_DIMENSIONS, (ULONG)NNUE_ACCUMULATOR_SIZE,
              (ULONG)NNUE_L2, (ULONG)NNUE_L3);
        fclose(fp);
        return FALSE;
    }

#define READ_OR_FAIL(ptr, count)                                        \
    do {                                                                \
        if ((size_t)(count) != fread((ptr), sizeof(float), (count), fp)) \
        {                                                               \
            Trace("NnueLoadWeights: %s truncated\n", szFilename);       \
            fclose(fp);                                                 \
            return FALSE;                                               \
        }                                                               \
    } while (0)

    READ_OR_FAIL(g_rgfNnueWeight, (size_t)NNUE_DIMENSIONS * NNUE_ACCUMULATOR_SIZE);
    READ_OR_FAIL(g_rgfNnueBias, NNUE_ACCUMULATOR_SIZE);
    READ_OR_FAIL(g_rgfNnueFc1Weight, (size_t)NNUE_L2 * 2 * NNUE_ACCUMULATOR_SIZE);
    READ_OR_FAIL(g_rgfNnueFc1Bias, NNUE_L2);
    READ_OR_FAIL(g_rgfNnueFc2Weight, (size_t)NNUE_L3 * NNUE_L2);
    READ_OR_FAIL(g_rgfNnueFc2Bias, NNUE_L3);
    READ_OR_FAIL(g_rgfNnueFc3Weight, NNUE_L3);
    READ_OR_FAIL(&g_rgfNnueFc3Bias, 1);

#undef READ_OR_FAIL

    fclose(fp);
    g_fNnueWeightsLoaded = TRUE;
    Trace("NnueLoadWeights: loaded %s\n", szFilename);
    return TRUE;
}


SCORE
NnueForward(SEARCHER_THREAD_CONTEXT *ctx)
/**

Routine description:

    Runs the full forward pass -- ReLU(accumulator), own-perspective-
    first concatenation by side to move (matching nnue_model.py's own
    convention so the single scalar output is side-to-move-relative
    by construction), then the dense 2*L1 -> L2 -> L3 -> 1 FC head --
    against ctx->sNnueAccum as it currently stands. Caller's
    responsibility to make sure that accumulator is actually correct
    for ctx->sPosition first (a fresh NnueAccumRecompute for both
    colors, if there's any doubt) -- this function trusts it as
    given.

Parameters:

    SEARCHER_THREAD_CONTEXT *ctx

Return value:

    SCORE : side-to-move-relative, in centipawns (the net's own raw
        output is in pawn units, same scale as the training label's
        score field -- see annotate_pgns/fit_k_and_label.py)

**/
{
    POSITION *pos = &ctx->sPosition;
    ULONG uStm = pos->uToMove;
    ULONG uOther = FLIP(uStm);
    float rgfConcat[2 * NNUE_ACCUMULATOR_SIZE];
    float rgfL2[NNUE_L2];
    float rgfL3[NNUE_L3];
    float fOut;
    ULONG i, j;

    for (i = 0; i < NNUE_ACCUMULATOR_SIZE; i++)
    {
        float v = ctx->sNnueAccum[uStm].rgfValue[i];
        rgfConcat[i] = (v > 0.0f) ? v : 0.0f;
    }
    for (i = 0; i < NNUE_ACCUMULATOR_SIZE; i++)
    {
        float v = ctx->sNnueAccum[uOther].rgfValue[i];
        rgfConcat[NNUE_ACCUMULATOR_SIZE + i] = (v > 0.0f) ? v : 0.0f;
    }

    for (j = 0; j < NNUE_L2; j++)
    {
        float s = g_rgfNnueFc1Bias[j];
        for (i = 0; i < 2 * NNUE_ACCUMULATOR_SIZE; i++)
        {
            s += g_rgfNnueFc1Weight[j][i] * rgfConcat[i];
        }
        rgfL2[j] = (s > 0.0f) ? s : 0.0f;
    }

    for (j = 0; j < NNUE_L3; j++)
    {
        float s = g_rgfNnueFc2Bias[j];
        for (i = 0; i < NNUE_L2; i++)
        {
            s += g_rgfNnueFc2Weight[j][i] * rgfL2[i];
        }
        rgfL3[j] = (s > 0.0f) ? s : 0.0f;
    }

    fOut = g_rgfNnueFc3Bias;
    for (i = 0; i < NNUE_L3; i++)
    {
        fOut += g_rgfNnueFc3Weight[i] * rgfL3[i];
    }

    return (SCORE)(fOut * 100.0f);
}