#!/usr/bin/env python3
"""Feature extraction driver: takes the mined FENs in mate_training_data.csv
and runs each through the engine's `evalkingdump` command (KING_DEFECT_DUMP
build only) to pull _EvalKing's raw per-defect breakdown, joining it back
onto each row's label/metadata to produce the final regression-ready CSV.

Sharded across N engine subprocesses (one `script` file each, run in
parallel) since 86k FENs through a single serial engine invocation would
be needlessly slow on this 24-core box -- same reasoning as
shard_mine.py's PGN sharding.

Each engine invocation processes its shard via `script <file>` containing
alternating `setboard <fen>` / `evalkingdump` lines; evalkingdump always
emits exactly two `KINGDEFECTDUMP color=...` lines (black then white, in
that fixed loop order) per position, unconditionally -- so the Nth pair
of KINGDEFECTDUMP lines in the shard's output corresponds to the Nth row
of that shard's input, and we pick out whichever color each row's
"mated_color_or_sampled_color" column names.
"""
import argparse
import csv
import os
import re
import subprocess
import sys
import time

KV_RE = re.compile(r"(\S+)=(\S+)")

# Fixed schema, since an off-board ring square (king near an edge/corner)
# emits only "sqN_offboard=1" while an on-board square emits 9 separate
# keys -- different king positions otherwise produce entirely different
# column sets, which would silently misalign rows if we inferred columns
# from whichever row happened to be parsed first (caught by a per-shard
# "feature key mismatch" warning during testing).
PIECE_NAMES = ("pawn", "minor", "rook", "queen")
PER_SQUARE_ON_BOARD_KEYS = [f"atk_{p}" for p in PIECE_NAMES] + \
    ["undef", "def", "xray", "flight"]
GLOBAL_KEYS = ["filestorm", "filestormopenfile", "filestormpawn",
               "queenproximity", "initloc", "countertotal",
               "safetyvalue", "attackerarmyscaler", "castled",
               "minorsathome", "misscastleshort", "misscastlelong"]

FEATURE_KEYS = []
for _sq in range(8):
    FEATURE_KEYS.append(f"sq{_sq}_offboard")
    for _k in PER_SQUARE_ON_BOARD_KEYS:
        FEATURE_KEYS.append(f"sq{_sq}_{_k}")
FEATURE_KEYS += GLOBAL_KEYS


def parse_kingdefectdump_line(line):
    """Returns (color, {key: value_str}) from a 'KINGDEFECTDUMP color=... ...'
    line, normalized to FEATURE_KEYS' fixed schema -- off-board squares get
    offboard=1 and all per-square on-board keys defaulted to 0 rather than
    left absent."""
    parts = line.strip().split(None, 1)
    assert parts[0] == "KINGDEFECTDUMP"
    raw_kv = dict(KV_RE.findall(parts[1]))
    color = raw_kv.pop("color")

    kv = {}
    for sq in range(8):
        offboard_key = f"sq{sq}_offboard"
        if offboard_key in raw_kv:
            kv[offboard_key] = raw_kv[offboard_key]
            for k in PER_SQUARE_ON_BOARD_KEYS:
                kv[f"sq{sq}_{k}"] = "0"
        else:
            kv[offboard_key] = "0"
            for k in PER_SQUARE_ON_BOARD_KEYS:
                full_key = f"sq{sq}_{k}"
                kv[full_key] = raw_kv[full_key]
    for k in GLOBAL_KEYS:
        kv[k] = raw_kv[k]
    assert set(kv.keys()) == set(FEATURE_KEYS), \
        f"parsed keys don't match fixed schema: {set(kv) ^ set(FEATURE_KEYS)}"
    return color, kv


def run_shard(engine_path, rows, shard_dir, shard_idx):
    script_path = os.path.join(shard_dir, f"extract_{shard_idx}.script")
    out_log = os.path.join(shard_dir, f"extract_{shard_idx}.out")
    with open(script_path, "w") as f:
        f.write("force\n")
        f.write("nopost\n")
        for row in rows:
            f.write(f"setboard {row['fen']}\n")
            f.write("evalkingdump\n")

    cmd = [engine_path, "--batch", "--command",
           f"script {script_path}"]
    with open(out_log, "w") as logf:
        rc = subprocess.run(cmd, stdout=logf, stderr=subprocess.STDOUT).returncode
    if rc != 0:
        raise RuntimeError(f"shard {shard_idx} engine exited {rc}, see {out_log}")

    dumps = []
    with open(out_log) as f:
        for line in f:
            if line.startswith("KINGDEFECTDUMP"):
                dumps.append(parse_kingdefectdump_line(line))

    if len(dumps) != 2 * len(rows):
        raise RuntimeError(
            f"shard {shard_idx}: expected {2 * len(rows)} KINGDEFECTDUMP "
            f"lines (2 per row), got {len(dumps)} -- a setboard likely "
            f"failed silently and evalkingdump re-used a stale position. "
            f"See {out_log}.")

    out_rows = []
    for i, row in enumerate(rows):
        black_color, black_kv = dumps[2 * i]
        white_color, white_kv = dumps[2 * i + 1]
        assert black_color == "black" and white_color == "white", \
            f"row {i}: unexpected color order {black_color}/{white_color}"
        target = row["mated_color_or_sampled_color"]
        kv = black_kv if target == "black" else white_kv
        out_row = dict(row)
        out_row.update(kv)
        out_rows.append(out_row)

    os.remove(script_path)
    os.remove(out_log)
    return out_rows


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("training_csv")
    ap.add_argument("out_csv")
    ap.add_argument("--engine", default="../typhoon_kingdump")
    ap.add_argument("--num-shards", type=int, default=8)
    ap.add_argument("--work-dir", default=None)
    args = ap.parse_args()

    engine_path = os.path.abspath(args.engine)
    if not os.path.isfile(engine_path):
        print(f"Engine not found: {engine_path}", file=sys.stderr)
        sys.exit(1)

    with open(args.training_csv, newline="") as f:
        rows = list(csv.DictReader(f))
    print(f"Loaded {len(rows)} rows from {args.training_csv}", file=sys.stderr)

    work_dir = args.work_dir or os.path.dirname(os.path.abspath(args.out_csv))
    os.makedirs(work_dir, exist_ok=True)

    n = args.num_shards
    shard_size = (len(rows) + n - 1) // n
    shards = [rows[i:i + shard_size] for i in range(0, len(rows), shard_size)]

    t0 = time.time()
    all_out_rows = [[] for _ in shards]

    import concurrent.futures
    with concurrent.futures.ThreadPoolExecutor(max_workers=n) as pool:
        futures = {}
        for idx, shard_rows in enumerate(shards):
            fut = pool.submit(run_shard, engine_path, shard_rows, work_dir, idx)
            futures[fut] = idx
        for fut in concurrent.futures.as_completed(futures):
            idx = futures[fut]
            out_rows = fut.result()
            all_out_rows[idx] = out_rows
            print(f"Shard {idx} done ({len(out_rows)} rows)", file=sys.stderr)

    print(f"All shards done in {time.time() - t0:.1f}s", file=sys.stderr)

    base_cols = ["fen", "label", "mated_color_or_sampled_color", "n_back",
                 "game_index", "control_source"]
    fieldnames = base_cols + FEATURE_KEYS

    with open(args.out_csv, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames)
        w.writeheader()
        for out_rows in all_out_rows:
            for row in out_rows:
                w.writerow({k: row.get(k) for k in fieldnames})

    total = sum(len(r) for r in all_out_rows)
    print(f"Wrote {total} feature rows to {args.out_csv}", file=sys.stderr)


if __name__ == "__main__":
    main()