import sys, os, struct
import numpy as np
import chess

# ---- HalfKP encoder (simplified but structurally compatible) ----
PIECE_TO_IDX = {chess.PAWN:0, chess.KNIGHT:1, chess.BISHOP:2, chess.ROOK:3, chess.QUEEN:4}
SQ_N   = 64
KING_N = 64
PS_N   = 5 * 64  # non-king piece types * squares
FEAT_N = 2 * KING_N * PS_N  # stm side + nstm side

def encode_halfkp(board: chess.Board):
    """Return a binary 0/1 HalfKP feature vector (np.int8, length FEAT_N), or None if kings missing."""
    k_stm = board.king(board.turn)
    k_nstm = board.king(not board.turn)
    if k_stm is None or k_nstm is None:
        return None
    x = np.zeros(FEAT_N, dtype=np.int8)

    def add_feat(k_sq, side_idx, piece, sq):
        if piece == chess.KING or piece not in PIECE_TO_IDX: return
        ps = PIECE_TO_IDX[piece]*64 + sq
        base = side_idx*KING_N*PS_N + k_sq*PS_N
        x[base + ps] = 1

    for color in [True, False]:
        for p in PIECE_TO_IDX:
            for sq in board.pieces(p, color):
                add_feat(k_stm, 0, p, sq)  # side=0 => stm features
                add_feat(k_nstm, 1, p, sq) # side=1 => nstm features
    return x

# ---- Binpack format (simple & fast) ----
# Header:
#   magic[8]    = b'BINPACK\0'
#   feats_i32   = FEAT_N
#   count_i64   = number of records
# Records (repeated 'count' times):
#   features: FEAT_N bytes (uint8 0/1)
#   target:   float32 (evaluation in pawns, stm POV)

MAGIC = b'BINPACK\x00'

def parse_line_to_board_and_eval(line):
    parts = line.strip().split()
    if len(parts) < 7:
        return None
    fen = " ".join(parts[:6])
    try:
        raw_eval = float(parts[-1])   # centipawns
    except:
        return None
    board = chess.Board(fen)
    # Stockfish scores are side-to-move POV; convert cp -> pawns
    target = raw_eval / 100.0
    return board, target

def convert_fen_to_binpack(in_path, out_path):
    # first pass: count valid lines to write header
    n = 0
    with open(in_path, "r") as f:
        for ln in f:
            x = parse_line_to_board_and_eval(ln)
            if not x: continue
            b, t = x
            feats = encode_halfkp(b)
            if feats is None: continue
            n += 1

    if n == 0:
        print("No valid positions found.")
        return

    with open(out_path, "wb") as out, open(in_path, "r") as f:
        # header
        out.write(MAGIC)
        out.write(struct.pack("<i", FEAT_N))
        out.write(struct.pack("<q", n))
        # data
        wrote = 0
        for ln in f:
            x = parse_line_to_board_and_eval(ln)
            if not x: continue
            b, target = x
            feats = encode_halfkp(b)
            if feats is None: continue
            out.write(feats.tobytes(order="C"))            # FEAT_N bytes
            out.write(struct.pack("<f", float(target)))    # 4 bytes
            wrote += 1
        assert wrote == n
    print(f"✅ Wrote {n} positions to {out_path}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python make_binpack_halfkp.py <input_fen+evals.txt> <output.binpack>")
        sys.exit(1)
    inp, outp = sys.argv[1], sys.argv[2]
    if not os.path.exists(inp):
        print(f"Input missing: {inp}"); sys.exit(1)
    convert_fen_to_binpack(inp, outp)