#!/usr/bin/env python3
"""Extracts the FeatureTransformer's psqtWeights table from a real
Stockfish .nnue file (nn-fdc1d0fe6455.nnue, added to Stockfish 2023-06-12,
rev 38e61663d836e062af0bc002814ad5149c4b7729 -- see
https://github.com/official-stockfish/networks).

Parsed against that exact commit's header-reading code (fetched and
read directly, not guessed/reverse-engineered by eye -- see
src/nnue/evaluate_nnue.cpp:94-125, src/nnue/nnue_feature_transformer.h:254-258,
src/nnue/nnue_architecture.h:44, src/nnue/features/half_ka_v2_hm.h):

Top-level header (evaluate_nnue.cpp read_header):
    version   uint32 LE
    hashValue uint32 LE
    size      uint32 LE  (length of description string)
    description  `size` bytes

FeatureTransformer section (Detail::read_parameters + its own
read_parameters):
    ft_header_hash  uint32 LE, must equal
                    HalfKAv2_hm::HashValue (0x7f234cb8) XOR (2*L1)
    biases          L1 int16 values
    weights         L1 * InputDimensions int16 values
    psqtWeights     PSQTBuckets * InputDimensions int32 values,
                    laid out psqtWeights[index*PSQTBuckets + k]

InputDimensions = SQUARE_NB*PS_NB/2 = 22528, IDENTICAL formula/indexing
to typhoon's own HalfKAv2_hm spec (docs/NNUE.md) -- same king buckets,
same orientation table, same piece-square planes. PSQTBuckets=8
(nnue_architecture.h). L1 for this net is NOT assumed -- derived below
from the stored ft hash, since Stockfish's own loader treats that
hash as the authoritative check, not file size or guessed offsets.

This script ONLY extracts psqtWeights (a literal per-feature-index
scalar linear value, one per material-count bucket) for comparison/
diagnostic purposes against our own training coverage gaps -- NOT for
copying into our net's weights. See the conversation: the goal is a
from-scratch net, this is a sanity check on whether our dusty corners
are in the right ballpark, not a transplant.

Usage:
    python3 parse_stockfish_net.py nn-fdc1d0fe6455.nnue --out sf_psqt.npy
"""
import argparse
import struct
import sys

import numpy as np

HALFKAV2_HM_HASH = 0x7F234CB8
PS_NB = 11 * 64  # 704
INPUT_DIMENSIONS = (64 * PS_NB) // 2  # 22528
PSQT_BUCKETS = 8


def read_u32(f):
    return struct.unpack("<I", f.read(4))[0]


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                  formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("nnue_file")
    ap.add_argument("--out", required=True,
                     help="output .npy, shape [INPUT_DIMENSIONS, PSQT_BUCKETS] int32")
    args = ap.parse_args()

    with open(args.nnue_file, "rb") as f:
        version = read_u32(f)
        top_hash = read_u32(f)
        desc_len = read_u32(f)
        desc = f.read(desc_len).decode("utf-8", errors="replace")
        print(f"version=0x{version:08x} top_hash=0x{top_hash:08x}")
        print(f"description: {desc!r}")

        ft_hash = read_u32(f)
        print(f"ft_hash=0x{ft_hash:08x}")

        # ft_hash = FeatureSet::HashValue ^ (OutputDimensions * 2)
        # (feature_transformer.h:250) -- solve for L1 (OutputDimensions).
        l1_times_2 = ft_hash ^ HALFKAV2_HM_HASH
        if l1_times_2 % 2 != 0:
            sys.exit(f"ft_hash XOR HalfKAv2_hm hash gave an odd value "
                      f"({l1_times_2}) -- this net's feature set isn't "
                      f"HalfKAv2_hm, or the hash constant is wrong for "
                      f"this Stockfish version. Refusing to guess further.")
        l1 = l1_times_2 // 2
        print(f"Derived L1 (TransformedFeatureDimensions) = {l1}")

        n_biases = l1
        n_weights = l1 * INPUT_DIMENSIONS
        n_psqt = PSQT_BUCKETS * INPUT_DIMENSIONS

        biases = np.frombuffer(f.read(2 * n_biases), dtype="<i2")
        assert len(biases) == n_biases, "short read on biases -- bad offset math"
        print(f"biases: read {len(biases)}, "
              f"range [{biases.min()}, {biases.max()}]")

        # Skip the weights array -- not needed for this diagnostic, and
        # it's ~66MB (l1 * 22528 * 2 bytes), the bulk of the file.
        f.seek(2 * n_weights, 1)

        psqt_raw = np.frombuffer(f.read(4 * n_psqt), dtype="<i4")
        if len(psqt_raw) != n_psqt:
            sys.exit(f"short read on psqtWeights: got {len(psqt_raw)}, "
                      f"expected {n_psqt} -- offset math above must be "
                      f"wrong (L1 derivation or file truncated)")

    psqt = psqt_raw.reshape(INPUT_DIMENSIONS, PSQT_BUCKETS)
    print(f"psqtWeights: shape {psqt.shape}, "
          f"range [{psqt.min()}, {psqt.max()}]")
    np.save(args.out, psqt)
    print(f"Wrote {args.out}")


if __name__ == "__main__":
    sys.exit(main())