src/chess_tutor/export_training_dataset.py

#!/usr/bin/env python3
"""
Export positions for ML training: Parquet (metadata + heuristics) + HDF5 (full ONNX tensors).

Reads rows from ``sampled_positions`` or ``evaluated_positions``. For each row, recomputes
the twelve heuristic scores from the FEN and runs ONNX once. Probe tensors default to
``tensors_to_expose`` in ``onnx_config.yaml`` (override with ``--probes``).

Stores:

  * ``<out_dir>/parquet/part-*.parquet`` — ``sample_idx``, ``db_id``, ``fen``, heuristics,
    and SQLite metadata (excludes ``onnx_eval_json``).
  * ``<out_dir>/tensors.h5`` — HDF5 datasets (float32), row-aligned by ``sample_idx``:
    ``policy``, ``wdl`` or ``value_head``, and one dataset per probe present in the ONNX
    outputs (names derived from ONNX tensor names; mapping in dataset attrs).

Example::

    python -m chess_tutor.export_training_dataset --table sampled_positions \\
        --out-dir training_export --limit 500000 --parquet-shard-rows 50000

Requires: onnxruntime, lczero.backends, numpy, pyarrow, h5py (see requirements.txt).
"""

from __future__ import annotations

import argparse
import json
import sqlite3
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import h5py
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

import lczero.backends as lc0

from chess_tutor.heuristics import HeuristicFeatures, extract_heuristic_features
from chess_tutor.lc0_onnx import (
    input_channels_from_onnx_shape,
    make_onnx_input_tensor,
    run_session,
)
from chess_tutor.compute import add_compute_arguments, log_compute_plan, resolve_lc0_bindings_backend
from chess_tutor.lc0_onnx.session import create_inference_session
from chess_tutor.onnx_probe_config import load_onnx_probe_config
from chess_tutor.paths import default_positions_db_path
from chess_tutor.training.heuristics_export import heuristics_row_dict
from chess_tutor.training.tensor_layout import (
    classify_outputs,
    merge_layout_batch,
    stack_batch,
)
from chess_tutor.training.tensor_names import tensor_dataset_name

# Backward-compatible re-exports (prefer ``chess_tutor.training``).
__all__ = [
    "ALLOWED_TABLES",
    "PARQUET_SKIP",
    "classify_outputs",
    "heuristics_row_dict",
    "iter_source_rows",
    "main",
    "merge_layout_batch",
    "sqlite_row_dict",
    "stack_batch",
    "table_column_names",
    "tensor_dataset_name",
    "write_parquet_shard",
]

ALLOWED_TABLES = frozenset({"sampled_positions", "evaluated_positions"})

PARQUET_SKIP = frozenset({"onnx_eval_json"})


def sqlite_row_dict(
    column_names: Sequence[str], row: Tuple[Any, ...]
) -> Dict[str, Any]:
    out: Dict[str, Any] = {}
    for name, val in zip(column_names, row):
        if name in PARQUET_SKIP:
            continue
        if name == "id":
            out["db_id"] = val
        else:
            out[name] = val
    return out


def table_column_names(conn: sqlite3.Connection, table: str) -> List[str]:
    rows = conn.execute(f"PRAGMA table_info({table})").fetchall()
    return [r[1] for r in rows]


def iter_source_rows(
    conn: sqlite3.Connection,
    table: str,
    limit: Optional[int],
) -> Tuple[List[str], Any]:
    """Return (column_names, cursor) for ``SELECT * FROM table ORDER BY id``."""
    names = table_column_names(conn, table)
    if not names:
        raise ValueError(f"Table {table!r} has no columns")
    col_list = ", ".join(f'"{n}"' for n in names)
    q = f"SELECT {col_list} FROM {table} ORDER BY id"
    params: Tuple[Any, ...] = ()
    if limit is not None:
        q += " LIMIT ?"
        params = (int(limit),)
    cur = conn.execute(q, params)
    return names, cur


def write_parquet_shard(rows: List[Dict[str, Any]], path: Path) -> None:
    table = pa.Table.from_pylist(rows)
    path.parent.mkdir(parents=True, exist_ok=True)
    pq.write_table(table, path, compression="zstd")


def main() -> None:
    cfg = load_onnx_probe_config()
    default_probe_names: Sequence[str] = list(cfg.tensors_to_expose)

    p = argparse.ArgumentParser(
        description="Export Parquet (metadata + heuristics) and HDF5 (full ONNX tensors)."
    )
    p.add_argument("--db", type=Path, default=default_positions_db_path())
    p.add_argument(
        "--table",
        default="sampled_positions",
        choices=sorted(ALLOWED_TABLES),
        help="SQLite table to read (must contain id, fen)",
    )
    p.add_argument(
        "--out-dir",
        type=Path,
        default=Path("training_export"),
        help="Output directory (creates parquet/ and tensors.h5)",
    )
    p.add_argument("--limit", type=int, default=None, help="Max rows (default: all)")
    p.add_argument(
        "--write-batch-rows",
        type=int,
        default=256,
        metavar="N",
        help="Accumulate N successful rows before appending to HDF5",
    )
    p.add_argument(
        "--parquet-shard-rows",
        type=int,
        default=50_000,
        metavar="N",
        help="Start a new Parquet file every N exported rows",
    )
    p.add_argument(
        "--hdf5-chunk-rows",
        type=int,
        default=8192,
        metavar="N",
        help="HDF5 chunk size along the sample (row) dimension",
    )
    p.add_argument("--onnx", type=Path, default=Path(cfg.onnx_path))
    p.add_argument("--weights", type=Path, default=Path(cfg.weights_path))
    p.add_argument(
        "--probes",
        nargs="*",
        default=None,
        help="Probe tensor names (default: tensors_to_expose from onnx_config.yaml)",
    )
    add_compute_arguments(p, include_ort=True, include_lc0_bindings=True)
    args = p.parse_args()
    log_compute_plan(ort_spec=args.ort_provider, lc0_bindings_spec=args.lc0_bindings_backend)
    probe_names: Sequence[str] = (
        args.probes if args.probes is not None else default_probe_names
    )

    if not args.db.is_file():
        print(f"Database not found: {args.db}", file=sys.stderr)
        sys.exit(1)
    if not args.onnx.is_file():
        print(f"ONNX not found: {args.onnx}", file=sys.stderr)
        sys.exit(1)
    if not args.weights.is_file():
        print(f"Weights not found: {args.weights}", file=sys.stderr)
        sys.exit(1)
    if args.write_batch_rows < 1 or args.parquet_shard_rows < 1 or args.hdf5_chunk_rows < 1:
        print("Batch and chunk sizes must be >= 1.", file=sys.stderr)
        sys.exit(1)

    args.out_dir.mkdir(parents=True, exist_ok=True)
    h5_path = args.out_dir / "tensors.h5"
    parquet_dir = args.out_dir / "parquet"
    parquet_dir.mkdir(parents=True, exist_ok=True)

    if h5_path.exists():
        print(
            f"Refusing to overwrite existing {h5_path}; remove it or pick a new --out-dir.",
            file=sys.stderr,
        )
        sys.exit(1)

    total_written = 0
    shard_idx = 0
    errors = 0
    h5_datasets: Dict[str, h5py.Dataset] = {}
    t0 = time.perf_counter()
    warned_missing_probes = False

    conn = sqlite3.connect(str(args.db))
    try:
        col_names, cursor = iter_source_rows(conn, args.table, args.limit)
        if "id" not in col_names or "fen" not in col_names:
            print(f"Table {args.table!r} must have id and fen columns.", file=sys.stderr)
            sys.exit(1)

        onnx_path = str(args.onnx.resolve())
        sess = create_inference_session(args.onnx.resolve(), ort_provider=args.ort_provider)
        onnx_in = sess.get_inputs()[0]
        if len(onnx_in.shape) != 4:
            raise ValueError(f"Unexpected ONNX input shape: {onnx_in.shape}")
        num_planes = input_channels_from_onnx_shape(onnx_in.shape)
        weights = lc0.Weights(str(args.weights.resolve()))
        backend = lc0.Backend(
            weights, backend=resolve_lc0_bindings_backend(args.lc0_bindings_backend)
        )

        out_names = {o.name for o in sess.get_outputs()}
        missing_cfg = [n for n in probe_names if n not in out_names]
        if missing_cfg:
            print(
                "Note: these names from --probes / onnx_config are not ONNX graph outputs "
                f"(often fixed by re-running the probe generator + reload): {missing_cfg}",
                file=sys.stderr,
            )

        h5f = h5py.File(h5_path, "w")
        try:
            parquet_buffer: List[Dict[str, Any]] = []
            pending_layout: List[Dict[str, Any]] = []
            reference_layout: Optional[Dict[str, Any]] = None

            def flush_h5_batch(layout_batch: List[Dict[str, Any]]) -> None:
                nonlocal h5_datasets, total_written
                stacked = stack_batch(layout_batch)
                b = stacked["policy"].shape[0]
                if not h5_datasets:
                    for key, arr in stacked.items():
                        tail = arr.shape[1:]
                        chunk_rows = max(1, min(args.hdf5_chunk_rows, b))
                        chunks = (chunk_rows,) + tail
                        h5_datasets[key] = h5f.create_dataset(
                            key,
                            shape=(0,) + tail,
                            maxshape=(None,) + tail,
                            dtype=np.float32,
                            chunks=chunks,
                            compression="gzip",
                            compression_opts=4,
                            shuffle=True,
                        )
                    probe_meta = [
                        {"onnx_name": p["onnx_name"], "dataset_key": p["dataset_key"]}
                        for p in layout_batch[0]["probes"]
                    ]
                    h5f.attrs["probes"] = json.dumps(probe_meta)
                    h5f.attrs["onnx_path"] = str(args.onnx.resolve())
                    h5f.attrs["sqlite_table"] = args.table
                    h5f.attrs["db_path"] = str(args.db.resolve())
                    h5f.attrs["tensor_datasets"] = json.dumps(list(stacked.keys()))

                pos = total_written
                new_len = pos + b
                for key, arr in stacked.items():
                    ds = h5_datasets[key]
                    ds.resize((new_len,) + ds.shape[1:])
                    ds[pos:new_len] = arr
                total_written = new_len

            def flush_parquet_shards(*, final: bool) -> None:
                nonlocal shard_idx, parquet_buffer
                if final:
                    while parquet_buffer:
                        n = min(args.parquet_shard_rows, len(parquet_buffer))
                        chunk = parquet_buffer[:n]
                        del parquet_buffer[:n]
                        out_p = parquet_dir / f"part-{shard_idx:05d}.parquet"
                        write_parquet_shard(chunk, out_p)
                        shard_idx += 1
                else:
                    while len(parquet_buffer) >= args.parquet_shard_rows:
                        chunk = parquet_buffer[: args.parquet_shard_rows]
                        del parquet_buffer[: args.parquet_shard_rows]
                        out_p = parquet_dir / f"part-{shard_idx:05d}.parquet"
                        write_parquet_shard(chunk, out_p)
                        shard_idx += 1

            for row in cursor:
                meta = sqlite_row_dict(col_names, row)
                db_id = meta.pop("db_id", None)
                fen = meta.get("fen")
                if db_id is None or not fen:
                    errors += 1
                    continue

                try:
                    h = extract_heuristic_features(str(fen))
                    hdict = heuristics_row_dict(h)
                except Exception:
                    errors += 1
                    continue

                try:
                    inp = make_onnx_input_tensor(str(fen), sess, backend, num_planes)
                    onnx_out = run_session(inp, sess)
                    layout = classify_outputs(sess, onnx_out, probe_names)
                    if not warned_missing_probes:
                        got = {p["onnx_name"] for p in layout["probes"]}
                        wanted = {n for n in probe_names if n in out_names}
                        skipped = sorted(wanted - got)
                        if skipped:
                            print(
                                "Note: configured probes not present in this run's ONNX outputs "
                                f"(skipped): {skipped}",
                                file=sys.stderr,
                            )
                        warned_missing_probes = True
                    reference_layout = merge_layout_batch(reference_layout, layout)
                except Exception:
                    errors += 1
                    continue

                sample_idx = total_written + len(pending_layout)
                pending_layout.append(layout)
                record: Dict[str, Any] = {
                    "sample_idx": sample_idx,
                    "db_id": db_id,
                    "fen": str(fen),
                    **hdict,
                    **{k: v for k, v in meta.items() if k != "fen"},
                }
                parquet_buffer.append(record)
                flush_parquet_shards(final=False)

                if len(pending_layout) >= args.write_batch_rows:
                    flush_h5_batch(pending_layout)
                    pending_layout.clear()

                n_live = total_written + len(pending_layout)
                if n_live % 5000 == 0 and n_live > 0:
                    dt = time.perf_counter() - t0
                    rate = n_live / dt if dt > 0 else 0
                    print(
                        f"  exported {n_live} rows  ({rate:.1f} rows/s)  errors={errors}",
                        flush=True,
                    )

            if pending_layout:
                flush_h5_batch(pending_layout)
                pending_layout.clear()

            flush_parquet_shards(final=True)

        finally:
            h5f.close()

    finally:
        conn.close()

    if total_written == 0:
        print(
            "No rows exported (all errors or empty table). Removing tensors.h5.",
            file=sys.stderr,
        )
        h5_path.unlink(missing_ok=True)
        sys.exit(1)

    h5_shapes: Dict[str, List[int]] = {}
    with h5py.File(h5_path, "r") as hf_r:
        for k in sorted(hf_r.keys()):
            h5_shapes[k] = list(hf_r[k].shape)

    manifest = {
        "exported_rows": total_written,
        "errors_skipped": errors,
        "parquet_shards": shard_idx,
        "parquet_dir": str(parquet_dir.relative_to(args.out_dir)),
        "hdf5_file": h5_path.name,
        "hdf5_tensor_shapes": h5_shapes,
        "elapsed_sec": round(time.perf_counter() - t0, 3),
        "write_batch_rows": args.write_batch_rows,
        "parquet_shard_rows": args.parquet_shard_rows,
        "hdf5_chunk_rows": args.hdf5_chunk_rows,
    }
    with open(args.out_dir / "export_manifest.json", "w", encoding="utf-8") as f:
        json.dump(manifest, f, indent=2)

    print(
        f"Done. exported={total_written} errors_skipped={errors} "
        f"out_dir={args.out_dir.resolve()}",
        flush=True,
    )


if __name__ == "__main__":
    main()