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

# ---- HalfKAv2_hm^ encoder ----
PIECE_TO_IDX = {
    chess.PAWN: 0,
    chess.KNIGHT: 1,
    chess.BISHOP: 2,
    chess.ROOK: 3,
    chess.QUEEN: 4,
    chess.KING: 5,  # king context channel
}

SQ_N = 64
KING_N = 64
PS_N = len(PIECE_TO_IDX) * SQ_N  # 6 * 64 = 384
FEAT_N = KING_N * PS_N  # 64 * 384 = 24576

MAGIC = b"BINPACK\x00"

def mirror_sq(sq: int) -> int:
    """Mirror square vertically (so black’s side looks like white’s)."""
    return chess.square_mirror(sq)

def encode_halfkav2_hm(board: chess.Board):
    """HalfKAv2_hm^ 0/1 feature vector for Stockfish 17.1 (length FEAT_N)."""
    k_sq = board.king(board.turn)
    if k_sq is None:
        return None
    x = np.zeros(FEAT_N, dtype=np.int8)

    # Mirror board if black to move so both sides share the same mapping
    def norm_sq(sq):
        return mirror_sq(sq) if board.turn == chess.BLACK else sq

    k_sq = norm_sq(k_sq)

    def add_feat(piece, sq):
        ps = PIECE_TO_IDX.get(piece)
        if ps is None: return
        sq = norm_sq(sq)
        idx = k_sq * PS_N + ps * SQ_N + sq
        x[idx] = 1

    for p in PIECE_TO_IDX:
        for color in [chess.WHITE, chess.BLACK]:
            for sq in board.pieces(p, color):
                add_feat(p, sq)

    return x

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
    b = chess.Board(fen)
    return b, raw_eval / 100.0  # convert to pawns

def convert_fen_to_binpack(in_path, out_path):
    # first pass: count valid positions
    n = 0
    with open(in_path) as f:
        for ln in f:
            x = parse_line_to_board_and_eval(ln)
            if not x: continue
            b, _ = x
            feats = encode_halfkav2_hm(b)
            if feats is not None:
                n += 1

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

    with open(out_path, "wb") as out, open(in_path) as f:
        out.write(MAGIC)
        out.write(struct.pack("<i", FEAT_N))
        out.write(struct.pack("<q", n))

        wrote = 0
        for ln in f:
            x = parse_line_to_board_and_eval(ln)
            if not x: continue
            b, t = x
            feats = encode_halfkav2_hm(b)
            if feats is None: continue
            out.write(feats.tobytes(order="C"))
            out.write(struct.pack("<f", float(t)))
            wrote += 1

        assert wrote == n
    print(f"Wrote {n} HalfKAv2_hm^ positions to {out_path}")

if __name__ == "__main__":
    if len(sys.argv) != 3:
        print("Usage: python make_binpack.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)