#!/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).

Skips exact rows (no p_search/r_outcome label -- see NNUE.md: a
mate/EGTB/draw-claim score is a game-theoretic fact, not something to
fit a sigmoid-WDL target against).

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
"""
import argparse
import array
import json
import sys
from pathlib import Path

import numpy as np

from nnue_feature_spec import active_feature_indices, chess


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 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=1.0,
                     help="target = lambda*p_search + (1-lambda)*r_outcome "
                          "(Phase A default 1.0, pure p_search -- see NNUE.md)")
    ap.add_argument("--limit", type=int, default=None,
                     help="stop after this many usable positions (smoke testing)")
    args = ap.parse_args()

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

    # 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 = []

    n_total = 0
    n_used = 0
    n_skipped = 0

    for pos in iter_positions(args.infile):
        n_total += 1
        if args.limit is not None and n_used >= args.limit:
            break
        if pos.get("exact") or "p_search" not in pos:
            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)
        p_search = pos["p_search"]
        r_outcome = pos.get("r_outcome", p_search)
        target.append(args.lam * p_search + (1.0 - args.lam) * r_outcome)

        n_used += 1
        if n_used % 200000 == 0:
            print(f"  ...{n_used} positions processed", file=sys.stderr)

    print(f"Read {n_total} rows ({args.infile})")
    print(f"  {n_skipped} skipped (exact / no p_search label)")
    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))

    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()