#!/usr/bin/env python3
"""The Phase A prototype net: a small, known HalfKAv2_hm architecture
per NNUE.md's architecture decisions -- 256x2-32-32-1-ish, a single
shared feature-transformer weight matrix applied independently to
each perspective's active-feature set (that sharing, not two
separate tables, is what makes the two 256-wide "accumulator" halves
comparable before concatenation), own-perspective-half-first
concatenation keyed by side-to-move so the net's single scalar
output is side-to-move-relative by construction (per the "Output
convention" architecture decision), and a sigmoid-WDL loss against
the already-fitted K=1.3123 (per "WDL-sigmoid target conversion
built and run", fit_k_and_label.py) rather than raw centipawn MSE.

Deliberately NOT doing yet (explicitly deferred, see NNUE.md): clipped
activations / quantization-aware scaling (this is a float32-only
Phase A prototype, no export-for-C-inference step exists yet), the
lazy-eval PSQT-term fast path, and the super-lazy material-only
prefilter. All irrelevant to proving the training pipeline works.

2026-10-17: added `material_probe`, a small auxiliary linear head on
the same concatenated accumulator `fc1` reads, trained jointly with
the main WDL loss against each position's raw material balance
(mover-relative pawns, no sigmoid). Added after a real, measured
failure: a net trained with plain MSE (and separately, exponent=2.6)
against WDL-sigmoid targets alone could not tell a lost pawn from a
lost queen in crowded midgame positions (see sanity_check_net.py) --
directionally correct but magnitude-blind.

2026-10-20: replaced by `score_probe` (same mechanism, different
target -- see that section below) and removed. Once `score_probe`
regresses Stockfish's actual raw score directly, `material_probe`'s
target (piece count) is a strict special case of what `score_probe`
already has to recover -- material dominates Stockfish's own score in
almost any position, so a net that fits the real score well has
implicitly recovered material too. Kept as one objective instead of
two overlapping ones; re-add material_probe alongside score_probe only
if a future sanity check shows score_probe alone is still
material-blind in some case it doesn't already cover.

2026-10-20: tried `score_probe`, same mechanism as material_probe
(small Linear head on the same accumulator, trained jointly, dropped
before export) targeting Stockfish's own raw per-position score
directly via plain MSE -- motivated by a measured failure of a
different kind from material_probe's: in clearly-losing positions,
the main WDL head's score magnitude did not reflect how lost the
position actually was (-2 for "lost" and -3 for "about to be mated",
no meaningful gradation), because the WDL loss only ever supervises
`raw` through `sigmoid(raw/K)` against a win-probability target, which
saturates hard once a position is clearly decided (K=2.0659 on the
2026-10-20 corpus -- sigmoid(8/K)=0.9796 vs sigmoid(13/K)=0.9982, a
target spread of under 0.02 for a 5-pawn difference in the real
score). **Measured, then reverted**: a medium-scale (500k positions,
~15 epochs) controlled comparison (score_probe-weighted run vs. an
otherwise-identical control) showed only a small improvement in how
well `forward()`'s own raw output tracked the real score at the
losing extremes -- compression was still clearly present in both.
Root cause: score_probe's gradient only ever reaches the *shared
accumulator* (`ft`/`ft_p`/`ft_bias`), never `fc1`/`fc2`/`fc3` --
nothing forces the actually-exported path to widen its own range, it
can only optionally make use of a richer upstream representation if
doing so happens to help the (still-saturating) WDL loss too.

2026-10-20: replaced score_probe with a direct, unsaturated regression
term on `raw` itself: `MSE(raw, clipped_score)`, added straight
alongside the WDL loss in nnue_train.py, no auxiliary head at all.
This puts gradient directly on `fc1`/`fc2`/`fc3` -- the actual
exported path -- rather than hoping a richer accumulator happens to
propagate through. Not a competing objective with the WDL loss, either:
since `p_search = sigmoid(score/K)`, the WDL loss's own implied ideal
`raw` is already approximately `score` (especially at lambda close to
1.0) -- this term is an unsaturated restatement of the same target,
not a different one, so no separate probe/model-surface-area is
needed to get its benefit. See nnue_train.py's --score-loss-weight.

2026-10-18: added `ft_p`, the actual Stockfish-doc "P" factor (see
NNUE.md's "feature factorization" status entry) -- a second,
king-bucket-INDEPENDENT EmbeddingBag over PS_NB=704 rows
(piece_square x piece_type x piece_color only), summed into the same
accumulator `ft` feeds. Motivation: HalfKAv2_hm's real 22528-row
table is king-bucket-conditioned, so a (piece, square, color) tuple
that's common in general chess but rare in a specific king bucket
gets its own mostly-untrained row -- measured directly on the
Stockfish self-play training pool: 9.9% of rows never fire at all,
23.6% see fewer than 50 examples. `ft_p`'s 704 rows see ~32x more
observations each (one per king bucket collapsed into one), supplying
dense gradient to a king-bucket-independent piece-value prior that
`ft`'s sparser rows can then specialize away from. Zero-initialized
(nn.init.zeros_) specifically so it contributes nothing at step 0 --
a default (unit-variance) init would silently double the
accumulator's effective scale at init, confounding the factorization
effect with an incidental learning-rate change. Deliberately no bias
of its own -- `ft_bias` stays the only bias on this sum, which is
also what keeps the post-training coalescing step (see
export_nnue_weights.py) exact rather than approximate. Training-only,
like material_probe: coalesced into `ft.weight` and dropped before
export, since `idx % PS_NB` (the P-factor's index for real feature
`idx`) is recoverable from the real index alone, needing no change to
build_training_tensors.py's cached tensors. Deliberately NOT adding
Stockfish's other two documented factors ("K", "HalfRelativeKP") --
HalfKAv2_hm already encodes the king directly via the shared PS_KING
plane, making a K factor redundant here; one variable at a time.
"""
import torch
import torch.nn as nn

from nnue_feature_spec import DIMENSIONS, PS_NB

L1 = 256   # feature-transformer output width (per perspective)
L2 = 32
L3 = 32

DEFAULT_K = 1.3123   # pawns; fit_k_and_label.py's Phase-A train-split fit


class NnueNet(nn.Module):
    def __init__(self, dimensions=DIMENSIONS, l1=L1, l2=L2, l3=L3,
                 ps_nb=PS_NB):
        super().__init__()
        # Single shared feature transformer: one [dimensions, l1]
        # weight matrix, applied once per perspective to that
        # perspective's own active-feature set. mode="sum" matches
        # Stockfish's own feature-transformer semantics (accumulate
        # one weight row per active feature).
        self.ft = nn.EmbeddingBag(dimensions, l1, mode="sum")
        self.ft_bias = nn.Parameter(torch.zeros(l1))

        # Training-only "P" factorization table -- see module
        # docstring's 2026-10-18 note. King-bucket-independent
        # (ps_nb=704 rows, vs ft's full 22528), summed into the same
        # accumulator. Zero-init so it contributes nothing at step 0.
        self.ft_p = nn.EmbeddingBag(ps_nb, l1, mode="sum")
        nn.init.zeros_(self.ft_p.weight)
        self.ps_nb = ps_nb

        self.fc1 = nn.Linear(2 * l1, l2)
        self.fc2 = nn.Linear(l2, l3)
        self.fc3 = nn.Linear(l3, 1)

    def accumulator(self, white_idx, white_off, black_idx, black_off, stm):
        """The ReLU'd, own-perspective-first concatenated accumulator
        (shape [B, 2*l1]) -- input to fc1. ft_p's
        contribution uses the identical bag offsets as ft -- same
        positions, same active-feature counts, just a second lookup
        of each index modulo ps_nb into the smaller, denser table."""
        white_idx_p = white_idx % self.ps_nb
        black_idx_p = black_idx % self.ps_nb

        white_acc = (self.ft(white_idx, white_off)
                     + self.ft_p(white_idx_p, white_off) + self.ft_bias)
        black_acc = (self.ft(black_idx, black_off)
                     + self.ft_p(black_idx_p, black_off) + self.ft_bias)

        white_acc = torch.relu(white_acc)
        black_acc = torch.relu(black_acc)

        stm = stm.view(-1, 1).float()
        # own-perspective-first concatenation, per side to move.
        own = stm * white_acc + (1.0 - stm) * black_acc
        enemy = stm * black_acc + (1.0 - stm) * white_acc
        return torch.cat([own, enemy], dim=1)

    def forward(self, white_idx, white_off, black_idx, black_off, stm):
        """stm: float/long tensor of shape [B], 1.0 if white to move
        else 0.0. Returns raw score in pawn-ish units, shape [B]."""
        x = self.accumulator(white_idx, white_off, black_idx, black_off, stm)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x.view(-1)

    def win_probability(self, white_idx, white_off, black_idx, black_off, stm, k=DEFAULT_K):
        raw = self.forward(white_idx, white_off, black_idx, black_off, stm)
        return torch.sigmoid(raw / k)