#!/usr/bin/env python3
"""Bit-for-bit cross-checker: diffs nnue_feature_spec.py's (Python)
HalfKAv2_hm feature indices against ../src/nnuefeatures.c's (C,
`dump nnuefeatures`) indices, for a set of hand-picked FENs chosen to
exercise every piece type, both king-file halves (a-d / e-h), both
perspectives, and non-symmetric positions (the starting position
alone can't catch a perspective-swap bug, since it's symmetric under
the exact transform between perspectives -- see NNUE.md).

Must pass before any training data is trusted to be feature-indexed
consistently between the Python trainer and the C inference code.

Usage:
    python3 cross_check_features.py [--engine ../src/typhoon]
"""
import argparse
import re
import subprocess
import sys
from pathlib import Path

from nnue_feature_spec import active_feature_indices, chess

NNUE_RE = re.compile(r"^NNUE_FEATURES (white|black) (.*)$")

# (label, FEN) -- deliberately includes:
#  - the symmetric starting position (sanity baseline only, see above)
#  - both kings on the queenside half (a-d files) and kingside half (e-h)
#  - a mixed case: White king queenside, Black king kingside (and the
#    reverse), to catch a perspective/orientation mixup
#  - every piece type present at least once, including promoted/extra
#    queens and multiple same-type pieces
#  - a near-empty king-and-pawns endgame (sparse feature set, edge of
#    the active-feature-count range)
TEST_POSITIONS = [
    ("startpos",
     "rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1"),
    ("both_kings_queenside",
     "2kr4/8/8/8/8/8/8/2KR4 w - - 0 1"),
    ("both_kings_kingside",
     "4rk2/8/8/8/8/8/8/4RK2 w - - 0 1"),
    ("white_queenside_black_kingside",
     "4rk2/8/8/8/8/8/8/2KR4 w - - 0 1"),
    ("white_kingside_black_queenside",
     "2kr4/8/8/8/8/8/8/4RK2 w - - 0 1"),
    ("midgame_mixed_material",
     "r1bqk2r/pp1nbppp/2p1pn2/3p4/2PP4/2N1PN2/PP2BPPP/R1BQK2R w KQkq - 0 8"),
    ("sparse_king_pawn_endgame",
     "8/5k2/8/3K4/8/8/5P2/8 w - - 0 1"),
    ("extra_queens",
     "3qk3/3q4/8/8/8/8/3Q4/3QK3 w - - 0 1"),
    ("black_to_move_midgame",
     "r2q1rk1/ppp2ppp/2n1bn2/2b1p3/4P3/2NPB3/PPP2PPP/R2Q1RK1 b - - 0 10"),
]


def run_engine_dump(engine, fen):
    cmd = [
        str(engine), "--cpus", "1", "--hash", "8m", "--logfile", "-",
        "--batch", "--command",
        f"force; setboard {fen}; dump nnuefeatures",
    ]
    out = subprocess.run(cmd, capture_output=True, text=True, timeout=30).stdout
    result = {}
    for line in out.splitlines():
        m = NNUE_RE.match(line.strip())
        if m:
            side, rest = m.groups()
            result[side] = sorted(int(x) for x in rest.split()) if rest else []
    return result


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--engine", default="../src/typhoon")
    args = ap.parse_args()

    engine = Path(args.engine).resolve()
    if not engine.exists():
        sys.exit(f"engine not found: {engine} (build it first -- see CLAUDE.md)")

    n_ok = 0
    n_fail = 0
    for label, fen in TEST_POSITIONS:
        board = chess.Board(fen)
        py = {
            "white": active_feature_indices(board, chess.WHITE),
            "black": active_feature_indices(board, chess.BLACK),
        }
        c_out = run_engine_dump(engine, fen)

        ok = True
        for side in ("white", "black"):
            if side not in c_out:
                print(f"[{label}] FAIL: engine produced no '{side}' line at all")
                ok = False
                continue
            if py[side] != c_out[side]:
                ok = False
                py_set, c_set = set(py[side]), set(c_out[side])
                print(f"[{label}] FAIL ({side}): "
                      f"python={len(py[side])} idx, C={len(c_out[side])} idx")
                only_py = sorted(py_set - c_set)
                only_c = sorted(c_set - py_set)
                if only_py:
                    print(f"    only in python: {only_py}")
                if only_c:
                    print(f"    only in C:      {only_c}")

        if ok:
            n_ok += 1
            print(f"[{label}] OK ({len(py['white'])} features/perspective)")
        else:
            n_fail += 1
            print(f"    FEN: {fen}")

    print(f"\n{n_ok} passed, {n_fail} failed, out of {len(TEST_POSITIONS)} positions")
    sys.exit(1 if n_fail else 0)


if __name__ == "__main__":
    main()