#!/usr/bin/env python3
"""Precompute HalfKAv2_hm feature tensors from a game-shaped labeled
file (quiet_filter_games.py output) into a cached, CSR-style on-disk
format that nnue_train.py can mmap and sample from directly --
avoids re-parsing FEN / recomputing SEE / recomputing feature indices
on every training epoch (the expensive part, done once here).

Each position's active feature indices (per nnue_feature_spec.py,
already cross-checked bit-for-bit against the C reference -- see
cross_check_features.py) are stored CSR-style (one flat concatenated
index array + a row-offset array per perspective), exactly the input
shape torch.nn.EmbeddingBag expects -- no fixed-width padding needed
since bag sizes vary (3-32 active features per position).

Exact rows (mate/EGTB/draw-claim positions) are never fit against the
sigmoid-WDL p_search target -- see NNUE.md: that score is a
game-theoretic fact, not noisy evidence to calibrate a sigmoid
against. But as of the EGTB-ground-truth-mixing decision, an exact row
that extraction could assign a direct r_outcome (1.0/0.0/0.5,
mover-relative) to is NOT dropped -- it is written to the tensors with
target = r_outcome directly, lambda-blend bypassed entirely, which is
exactly what the "mix TB ground truth into training" literature
describes. Only rows with neither p_search nor a determinable
r_outcome (unparseable comments, the "only"-move placeholder, an
exact row from data predating this fix) are still skipped.

Usage:
    python3 build_training_tensors.py \
        --in runs/sd10_twic_2M/games.train.labeled.quiet.jsonl \
        --outdir runs/sd10_twic_2M/tensors_train \
        --lambda 1.0

Per-position feature-index computation (chess.Board(fen) + SEE-free
active_feature_indices) is the expensive, embarrassingly-parallel part
-- positions have no cross-dependency. --workers (default 16) shards
the input across that many processes via multiprocessing.Pool,
stitching results back together in input order so output is
byte-for-byte identical to the old single-threaded path. --workers 1
restores that original sequential path; --limit (smoke testing) always
forces --workers 1 since it depends on exact file-order position
counting.
"""
import argparse
import array
import json
import sys
from multiprocessing import Pool
from pathlib import Path

import numpy as np

from nnue_feature_spec import active_feature_indices, chess

# Clip bound for the auxiliary score-probe target (nnue_model.py's
# score_probe, added 2026-10-20 -- see that module's docstring). A
# handful of blowout-but-not-forced-mate scores (see the user's own
# +13.67-before-mate example) would otherwise dominate the probe's
# MSE; exact/mate rows are excluded entirely (see SCORE_VALID below),
# not just clipped, since a mate score isn't "the Stockfish score" in
# the pawn sense this probe targets.
SCORE_CLIP_PAWNS = 20.0


def position_target(pos, lam):
    """Returns the training target for one position, or None if it's
    unusable (no p_search and no determinable r_outcome). Exact rows
    with a direct r_outcome bypass p_search/lambda entirely -- see
    this module's docstring and NNUE.md's EGTB-ground-truth-mixing
    decision."""
    if pos.get("exact"):
        return pos.get("r_outcome")
    if "p_search" not in pos:
        return None
    p_search = pos["p_search"]
    r_outcome = pos.get("r_outcome", p_search)
    return lam * p_search + (1.0 - lam) * r_outcome


def score_probe_target(pos):
    """Returns (score, valid) for the auxiliary score-probe target --
    Stockfish's own raw per-position score (pos["score"], already
    mover-relative pawns per extract_selfplay_positions.py), clipped
    to +-SCORE_CLIP_PAWNS, or (0.0, False) for exact/mate rows (no
    "score" field, or not a real Stockfish-judged centipawn-ish value
    -- same exclusion reasoning as p_search's, see NNUE.md)."""
    if pos.get("exact") or "score" not in pos:
        return 0.0, False
    s = pos["score"]
    s = max(-SCORE_CLIP_PAWNS, min(SCORE_CLIP_PAWNS, s))
    return s, True


def iter_positions(path):
    with open(path) as f:
        for line in f:
            game = json.loads(line)
            for pos in game["positions"]:
                yield pos


def iter_chunks(iterable, chunk_size):
    """Batches an iterable into lists of up to chunk_size items -- the
    unit of work handed to one worker process at a time."""
    chunk = []
    for item in iterable:
        chunk.append(item)
        if len(chunk) >= chunk_size:
            yield chunk
            chunk = []
    if chunk:
        yield chunk


def process_chunk(args):
    """Runs in a worker process: the same per-position feature-index /
    target / material computation main() used to do inline, just over
    one chunk. Returns flat, already-concatenated arrays for the chunk
    (not yet offset-adjusted -- that happens once, cheaply, back in
    the main process when chunks are stitched together in order) plus
    a per-position length list so the caller can rebuild the CSR
    offset array without re-deriving it from the index array."""
    positions, lam = args
    white_idx = array.array("i")
    black_idx = array.array("i")
    white_lens = array.array("i")
    black_lens = array.array("i")
    stm = []
    target = []
    score_probe_val = []
    score_probe_valid = []
    n_total = len(positions)
    n_skipped = 0

    for pos in positions:
        t = position_target(pos, lam)
        if t is None:
            n_skipped += 1
            continue

        board = chess.Board(pos["fen"])
        w = active_feature_indices(board, chess.WHITE)
        b = active_feature_indices(board, chess.BLACK)

        white_idx.extend(w)
        white_lens.append(len(w))
        black_idx.extend(b)
        black_lens.append(len(b))

        stm.append(1 if board.turn == chess.WHITE else 0)
        target.append(t)
        s, valid = score_probe_target(pos)
        score_probe_val.append(s)
        score_probe_valid.append(valid)

    return (white_idx, white_lens, black_idx, black_lens, stm, target,
            score_probe_val, score_probe_valid, n_total, n_skipped)


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--in", dest="infile", required=True)
    ap.add_argument("--outdir", required=True)
    ap.add_argument("--lambda", dest="lam", type=float, default=0.9,
                     help="target = lambda*p_search + (1-lambda)*r_outcome "
                          "for non-exact rows (exact rows always use "
                          "r_outcome directly, ignoring lambda). Default "
                          "0.9: with a Stockfish-sourced self-play corpus, "
                          "p_search (SF's own WDL-calibrated search score) "
                          "is a much lower-variance per-position signal "
                          "than one game's actual outcome, so most of the "
                          "weight stays on it -- see NNUE.md)")
    ap.add_argument("--limit", type=int, default=None,
                     help="stop after this many usable positions (smoke "
                          "testing) -- forces --workers 1, since the limit "
                          "is applied to the global used-position count in "
                          "file order, which only a single sequential pass "
                          "can honor exactly")
    ap.add_argument("--workers", type=int, default=16,
                     help="worker processes for the per-position feature-"
                          "index computation (the expensive, embarrassingly"
                          "-parallel part -- each position is independent). "
                          "1 disables multiprocessing and runs the original "
                          "sequential path")
    ap.add_argument("--chunk-size", type=int, default=20000,
                     help="positions per unit of work handed to a worker; "
                          "only relevant when --workers > 1")
    args = ap.parse_args()

    outdir = Path(args.outdir)
    outdir.mkdir(parents=True, exist_ok=True)

    workers = args.workers
    if args.limit is not None and workers != 1:
        print("--limit given: forcing --workers 1 for exact, "
              "file-order-respecting smoke-test behavior", file=sys.stderr)
        workers = 1

    # array.array('i') is a packed C int buffer -- avoids the ~28
    # bytes/element overhead of a Python list of boxed ints, which
    # matters once nnz reaches tens of millions.
    white_idx = array.array("i")
    black_idx = array.array("i")
    white_off = [0]
    black_off = [0]
    stm = []       # 1 = white to move, 0 = black
    target = []
    score_probe_val = []    # auxiliary score-probe target, clipped pawns
    score_probe_valid = []  # False for exact/mate rows -- excluded from that loss

    n_total = 0
    n_used = 0
    n_skipped = 0

    if workers == 1:
        for pos in iter_positions(args.infile):
            n_total += 1
            if args.limit is not None and n_used >= args.limit:
                break
            t = position_target(pos, args.lam)
            if t is None:
                n_skipped += 1
                continue

            board = chess.Board(pos["fen"])
            w = active_feature_indices(board, chess.WHITE)
            b = active_feature_indices(board, chess.BLACK)

            white_idx.extend(w)
            white_off.append(white_off[-1] + len(w))
            black_idx.extend(b)
            black_off.append(black_off[-1] + len(b))

            stm.append(1 if board.turn == chess.WHITE else 0)
            target.append(t)
            s, valid = score_probe_target(pos)
            score_probe_val.append(s)
            score_probe_valid.append(valid)

            n_used += 1
            if n_used % 200000 == 0:
                print(f"  ...{n_used} positions processed", file=sys.stderr)
    else:
        # Chunked across worker processes, but stitched back together
        # in input order (Pool.imap preserves task order, unlike
        # imap_unordered) -- output is byte-for-byte identical to the
        # sequential path, just computed faster on an idle multi-core
        # box. Each chunk's own CSR lengths get turned into this
        # process's running cumulative offsets here, cheaply, since
        # that's the one step that has to stay sequential.
        chunks = ((chunk, args.lam)
                  for chunk in iter_chunks(iter_positions(args.infile),
                                            args.chunk_size))
        with Pool(processes=workers) as pool:
            for (w_idx, w_lens, b_idx, b_lens, chunk_stm, chunk_target,
                 chunk_score_val, chunk_score_valid, chunk_total,
                 chunk_skipped) in pool.imap(process_chunk, chunks):
                n_total += chunk_total
                n_skipped += chunk_skipped

                white_idx.extend(w_idx)
                for length in w_lens:
                    white_off.append(white_off[-1] + length)
                black_idx.extend(b_idx)
                for length in b_lens:
                    black_off.append(black_off[-1] + length)

                stm.extend(chunk_stm)
                target.extend(chunk_target)
                score_probe_val.extend(chunk_score_val)
                score_probe_valid.extend(chunk_score_valid)

                n_used += len(chunk_stm)
                print(f"  ...{n_used} positions processed "
                      f"({n_total} read)", file=sys.stderr)

    print(f"Read {n_total} rows ({args.infile})")
    print(f"  {n_skipped} skipped (no p_search and no determinable r_outcome)")
    print(f"  {n_used} usable positions")

    np.save(outdir / "white_idx.npy", np.frombuffer(white_idx, dtype=np.int32))
    np.save(outdir / "white_off.npy", np.array(white_off, dtype=np.int64))
    np.save(outdir / "black_idx.npy", np.frombuffer(black_idx, dtype=np.int32))
    np.save(outdir / "black_off.npy", np.array(black_off, dtype=np.int64))
    np.save(outdir / "stm.npy", np.array(stm, dtype=np.uint8))
    np.save(outdir / "target.npy", np.array(target, dtype=np.float32))
    np.save(outdir / "score_probe_val.npy",
            np.array(score_probe_val, dtype=np.float32))
    np.save(outdir / "score_probe_valid.npy",
            np.array(score_probe_valid, dtype=np.uint8))

    nnz_white = len(white_idx)
    nnz_black = len(black_idx)
    print(f"Wrote tensors to {outdir}: {n_used} positions, "
          f"{nnz_white} white / {nnz_black} black nonzero feature entries "
          f"({nnz_white / max(n_used, 1):.1f} avg features/position)")


if __name__ == "__main__":
    main()