chess-engine/cleaning/format.py

#!/usr/bin/env python3
"""
Turn refined PGN games (Site/WhiteElo/BlackElo headers + movetext with
%eval comments) into a CSV where each row is the board state *after* one
move.

Row layout (67 columns):
    game_id, a8, b8, c8, d8, e8, f8, g8, h8, a7, ... , h1, eval, turn

  - game_id : 0-based index of the game in the input file. Not asked for
              explicitly, but with many games concatenated into one CSV
              there is no other way to tell which rows belong to which
              game, so I added it as column 1. Drop it with --no-game-id
              if you don't want it.
  - a8..h1  : the 64 squares (rank 8 down to rank 1, file a to h), each
              holding the FEN piece letter on that square after the move
              (uppercase = white, lowercase = black) or an empty string
              if the square is empty.
  - eval    : the %eval attached to that move (blank if that particular
              move had none).
  - turn    : "white" or "black" -- whose turn it is to move NEXT.

Streams game-by-game (like the two earlier scripts) so it stays cheap on
RAM no matter how large the input file is. Requires `pip install chess`.

Usage:
    python pgn_to_csv.py refined.pgn out.csv
    python pgn_to_csv.py refined.pgn out.csv --no-game-id
"""

import argparse
import csv
import re
import sys

import chess

TAG_RE = re.compile(r'^\[(\w+)\s+"(.*)"\]\s*$')
EVAL_RE = re.compile(r'\[%eval\s+([^\]]+)\]')
MOVE_NUM_RE = re.compile(r'^\d+\.+$')
RESULT_RE = re.compile(r'^(1-0|0-1|1/2-1/2|\*)$')
# Strip trailing annotation glyphs python-chess's SAN parser won't accept
# (e.g. "e5?!", "Nc3!!", "Qxd1+!") while keeping the "+"/"#" that ARE part
# of legal SAN.
SUFFIX_ANNOTATION_RE = re.compile(r'[!?]+$')

SQUARE_COLUMNS = [f + str(r) for r in range(8, 0, -1) for f in "abcdefgh"]


def tokenize_movetext(movetext: str):
    """Yield (san_move, eval_str_or_None) pairs from a movetext blob."""
    tokens = re.findall(r'\{[^}]*\}|\S+', movetext)
    i = 0
    n = len(tokens)
    while i < n:
        tok = tokens[i]
        if MOVE_NUM_RE.match(tok) or RESULT_RE.match(tok):
            i += 1
            continue
        move_san = tok
        eval_val = None
        if i + 1 < n and tokens[i + 1].startswith('{'):
            m = EVAL_RE.search(tokens[i + 1])
            if m:
                eval_val = m.group(1)
            i += 2
        else:
            i += 1
        yield move_san, eval_val


def board_row(board: chess.Board):
    """Return the 64 square values for the current board, a8..h1 order."""
    row = []
    for name in SQUARE_COLUMNS:
        sq = chess.parse_square(name)
        piece = board.piece_at(sq)
        row.append(piece.symbol() if piece else "")
    return row


def process_game(game_id, header_lines, movetext_lines, writer, include_game_id):
    if not movetext_lines:
        return 0

    movetext = ' '.join(line.strip() for line in movetext_lines if line.strip())

    board = chess.Board()
    rows_written = 0

    for move_san, eval_val in tokenize_movetext(movetext):
        clean_san = SUFFIX_ANNOTATION_RE.sub('', move_san)
        try:
            board.push_san(clean_san)
        except (chess.IllegalMoveError, chess.InvalidMoveError, chess.AmbiguousMoveError, ValueError) as e:
            print(f"  [game {game_id}] stopping early: could not parse move "
                  f"'{move_san}' ({e})", file=sys.stderr)
            break

        turn = "white" if board.turn == chess.WHITE else "black"
        row = (board_row(board) + [eval_val if eval_val is not None else "", turn])
        if include_game_id:
            row = [game_id] + row
        writer.writerow(row)
        rows_written += 1

    return rows_written


def refine(input_path: str, output_path: str, include_game_id: bool):
    header_lines = []
    movetext_lines = []
    state = "seek_header"
    game_id = 0
    total_rows = 0

    header_row = SQUARE_COLUMNS + ["eval", "turn"]
    if include_game_id:
        header_row = ["game_id"] + header_row

    with open(input_path, "r", encoding="utf-8", errors="replace", buffering=1024 * 1024) as fin, \
         open(output_path, "w", encoding="utf-8", newline="") as fout:

        writer = csv.writer(fout)
        writer.writerow(header_row)

        for raw_line in fin:
            line = raw_line.rstrip("\n")

            if state == "seek_header":
                if line.startswith("["):
                    header_lines.append(line)
                elif line.strip() == "" and header_lines:
                    state = "seek_movetext"
            elif state == "seek_movetext":
                if line.strip() == "":
                    total_rows += process_game(game_id, header_lines, movetext_lines,
                                                writer, include_game_id)
                    game_id += 1
                    header_lines = []
                    movetext_lines = []
                    state = "seek_header"
                else:
                    movetext_lines.append(line)

        if header_lines and movetext_lines:
            total_rows += process_game(game_id, header_lines, movetext_lines,
                                        writer, include_game_id)
            game_id += 1

    return game_id, total_rows


def main():
    parser = argparse.ArgumentParser(
        description="Convert refined PGN games into a per-move board-state CSV.")
    parser.add_argument("input", help="Path to the refined PGN file")
    parser.add_argument("output", help="Path to write the CSV file")
    parser.add_argument("--no-game-id", action="store_true",
                         help="Omit the leading game_id column (66-col strict layout).")
    args = parser.parse_args()

    games, rows = refine(args.input, args.output, include_game_id=not args.no_game_id)
    print(f"Processed {games} games -> {rows} rows written to {args.output}")


if __name__ == "__main__":
    main()