#!/usr/bin/env python3
"""First-principles pipeline sanity check, 2026-10-17: can the net
even OVERFIT a small, fixed, noiseless target -- the engine's own
classic Eval() score on a few thousand real positions -- regressed
directly (plain MSE on raw pawn-ish units, no WDL sigmoid, no K, no
lambda, no quiet filter, no quantization)?

Motivation: five independent head-to-head matches today (two nets,
two time controls, with and without the lazy probe) all showed NNUE
losing decisively to classic HCE, including at matched search depth
where ECM tactical-suite results had suggested NNUE should do well.
Before concluding anything about net/data quality, rule out a
structural pipeline bug the hard way: if this architecture, this
feature-indexing, and this training loop cannot drive error to
~zero on a target it should be trivially able to memorize (a few
thousand noiseless (position, scalar) pairs, same units/sign
convention as everything else in this pipeline), something is
broken in feature indexing, accumulator math, or the training loop
itself -- not a "real chess judgment is hard" problem.

If this DOES overfit cleanly, the pipeline is vindicated at the
mechanical level, and today's match losses are a real net-quality
question, not a bug hunt.

Usage:
    python3 overfit_sanity_check.py --tensor-dir /tmp/overfit_tensors \
        --epochs 400
"""
import argparse
import time
from pathlib import Path

import numpy as np
import torch

from nnue_model import NnueNet


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                  formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--tensor-dir", required=True)
    ap.add_argument("--epochs", type=int, default=400)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    torch.manual_seed(args.seed)
    d = Path(args.tensor_dir)
    # nn.EmbeddingBag's `offsets` wants one START index per bag (length
    # N) -- the CSR files store N+1 entries (with a trailing end
    # sentinel, per build_training_tensors.py), so the last entry
    # must be dropped here (nnue_train.py's _gather_batch does the
    # equivalent by only ever building N-length offset arrays).
    white_idx = torch.from_numpy(np.load(d / "white_idx.npy").astype(np.int64))
    white_off = torch.from_numpy(np.load(d / "white_off.npy")[:-1])
    black_idx = torch.from_numpy(np.load(d / "black_idx.npy").astype(np.int64))
    black_off = torch.from_numpy(np.load(d / "black_off.npy")[:-1])
    stm = torch.from_numpy(np.load(d / "stm.npy").astype(np.int64))
    target = torch.from_numpy(np.load(d / "target_raw.npy"))
    n = len(target)
    print(f"Loaded {n} (position, classic_eval) pairs -- target range "
          f"[{target.min():.2f}, {target.max():.2f}] pawns")

    model = NnueNet()
    optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)

    t0 = time.time()
    for epoch in range(args.epochs):
        optimizer.zero_grad()
        pred = model.forward(white_idx, white_off, black_idx, black_off, stm)
        loss = torch.nn.functional.mse_loss(pred, target)
        loss.backward()
        optimizer.step()
        if epoch % 20 == 0 or epoch == args.epochs - 1:
            rmse_cp = (loss.item() ** 0.5) * 100
            print(f"  epoch {epoch:4d}  train MSE (pawns^2) = {loss.item():.5f}  "
                  f"RMSE = {rmse_cp:.1f}cp  ({time.time() - t0:.0f}s)")

    with torch.no_grad():
        pred = model.forward(white_idx, white_off, black_idx, black_off, stm)
    final_mse = torch.nn.functional.mse_loss(pred, target).item()
    final_rmse_cp = (final_mse ** 0.5) * 100
    max_err_cp = (pred - target).abs().max().item() * 100
    print(f"\nFinal: train MSE = {final_mse:.5f} pawns^2, "
          f"RMSE = {final_rmse_cp:.1f}cp, max |error| = {max_err_cp:.1f}cp")
    print("Verdict: ", end="")
    if final_rmse_cp < 10:
        print("CLEAN OVERFIT -- pipeline mechanically sound. "
              "Today's match losses are a net-quality/data question, not a bug.")
    elif final_rmse_cp < 50:
        print("PARTIAL overfit -- some signal getting through, but not "
              "cleanly memorizing a noiseless target. Worth investigating "
              "further (LR, architecture capacity, or a real bug) before "
              "trusting any training result built on this pipeline.")
    else:
        print("FAILED to overfit even a tiny noiseless target -- strong "
              "evidence of a structural bug (feature indexing, accumulator "
              "math, or training loop), not a net-quality/data problem. "
              "Do not trust any prior training result until this is found.")


if __name__ == "__main__":
    main()