CPG/src/Chess/compare_positions.py

import re
import sys
import time
import subprocess
from shutil import which
from pathlib import Path
from argparse import ArgumentParser, Namespace



def main(args: Namespace) -> int:
    print("Nimfish move validator v0.0.1 by nocturn9x")
    try:
        STOCKFISH = (args.stockfish or Path(which("stockfish"))).resolve(strict=True)
    except Exception as e:
        print(f"Could not locate stockfish executable -> {type(e).__name__}: {e}")
        return -1
    try:
        NIMFISH = (args.nimfish or (Path.cwd() / "nimfish")).resolve(strict=True)
    except Exception as e:
        print(f"Could not locate nimfish executable -> {type(e).__name__}: {e}")
        return -1
    print(f"Starting Stockfish engine at {STOCKFISH.as_posix()!r}")
    stockfish_process = subprocess.Popen(STOCKFISH,
                                         stdout=subprocess.PIPE,
                                         stderr=subprocess.STDOUT,
                                         stdin=subprocess.PIPE,
                                         encoding="u8"
                                         )
    print(f"Starting Nimfish engine at {NIMFISH.as_posix()!r}")
    nimfish_process = subprocess.Popen(NIMFISH,
                                       stdout=subprocess.PIPE,
                                       stderr=subprocess.STDOUT,
                                       stdin=subprocess.PIPE,
                                       encoding="u8")
    print(f"Setting position to {(args.fen if args.fen else 'rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1')!r}")
    if args.fen:
        nimfish_process.stdin.write(f"position fen {args.fen}\n")
        stockfish_process.stdin.write(f"position fen {args.fen}\n")
    else:
        nimfish_process.stdin.write("position startpos\n")
        stockfish_process.stdin.write("position startpos\n")
    print(f"Engines started, beginning search to depth {args.ply}")
    nimfish_process.stdin.write(f"go perft {args.ply} {'bulk' if args.bulk else ''}\n")
    stockfish_process.stdin.write(f"go perft {args.ply}\n")
    print("Search started, waiting for engine completion")
    start_time = time.time()
    stockfish_output = stockfish_process.communicate()[0]
    stockfish_time = time.time() - start_time
    start_time = time.time()
    nimfish_output = nimfish_process.communicate()[0]
    nimfish_time = time.time() - start_time
    positions = {
        "all": {},
        "stockfish": {},
        "nimfish": {}
    }
    pattern = re.compile(r"(?P<source>[a-h][1-8])(?P<target>[a-h][1-8])(?P<promotion>b|n|q|r)?:\s(?P<nodes>[0-9]+)", re.MULTILINE)
    for (source, target, promotion, nodes) in pattern.findall(stockfish_output):
        move = f"{source}{target}{promotion}"
        positions["all"][move] = [int(nodes)]
        positions["stockfish"][move] = int(nodes)
    for (source, target, promotion, nodes) in pattern.findall(nimfish_output):
        move = f"{source}{target}{promotion}"
        if move in positions["all"]:
            positions["all"][move].append(int(nodes))
        else:
            positions["all"][move] = [int(nodes)]
        positions["nimfish"][move] = int(nodes)
    
    missing = {
        # Are in nimfish but not in stockfish
        "nimfish": [],
        # Are in stockfish but not in nimfish
        "stockfish": []
    }
    # What mistakes did Nimfish do?
    mistakes = set()
    for move, nodes in positions["all"].items():
        if move not in positions["stockfish"]:
            missing["nimfish"].append(move)
            continue
        elif move not in positions["nimfish"]:
            missing["stockfish"].append(move)
            continue
        if nodes[0] != nodes[1]:
            mistakes.add(move)
    
    total_nodes = {"stockfish": sum(positions["stockfish"][move] for move in positions["stockfish"]),
                   "nimfish": sum(positions["nimfish"][move] for move in positions["nimfish"])}
    total_difference = total_nodes["stockfish"] - total_nodes["nimfish"]
    print(f"Stockfish searched {total_nodes['stockfish']} node{'' if total_nodes['stockfish'] == 1 else 's'} in {stockfish_time:.2f} seconds")
    print(f"Nimfish searched {total_nodes['nimfish']} node{'' if total_nodes['nimfish'] == 1 else 's'} in {nimfish_time:.2f} seconds")

    if total_difference > 0:
        print(f"Nimfish searched {total_difference} less node{'' if total_difference == 1 else 's'} than Stockfish")
    elif total_difference < 0:
        total_difference = abs(total_difference)
        print(f"Nimfish searched {total_difference} more node{'' if total_difference == 1 else 's'} than Stockfish")
    else:
        print("Node count is identical")
    pattern = re.compile(r"(?:\s\s-\sCaptures:\s(?P<captures>[0-9]+))\n"
                            r"(?:\s\s-\sChecks:\s(?P<checks>[0-9]+))\n"
                            r"(?:\s\s-\sE\.P:\s(?P<enPassant>[0-9]+))\n"
                            r"(?:\s\s-\sCheckmates:\s(?P<checkmates>[0-9]+))\n"
                            r"(?:\s\s-\sCastles:\s(?P<castles>[0-9]+))\n"
                            r"(?:\s\s-\sPromotions:\s(?P<promotions>[0-9]+))",
                            re.MULTILINE)
    extra: re.Match | None = None
    if not args.bulk:
        extra = pattern.search(nimfish_output)
    missed_total = len(missing['stockfish']) + len(missing['nimfish'])
    if missing["stockfish"] or missing["nimfish"] or mistakes:
        print(f"Found {missed_total} missed move{'' if missed_total == 1 else 's'} and {len(mistakes)} counting mistake{'' if len(mistakes) == 1 else 's'}, more info below: ")
        if args.bulk:
            print("Note: Nimfish was run in bulk-counting mode, so a detailed breakdown of each move type is not available. "
                  "To fix this, re-run the program without the --bulk option")
        if extra:
            print(f"   Breakdown by move type:")
            print(f"    - Captures: {extra.group('captures')}")
            print(f"    - Checks: {extra.group('checks')}")
            print(f"    - En Passant: {extra.group('enPassant')}")
            print(f"    - Checkmates: {extra.group('checkmates')}")
            print(f"    - Castles: {extra.group('castles')}")
            print(f"    - Promotions: {extra.group('promotions')}")
            print(f"    - Total: {total_nodes['nimfish']}")

        elif not args.bulk:
            print("Unable to locate move breakdown in Nimfish output")
        if missing["stockfish"] or missing["nimfish"]:
            print("\n   Move count breakdown:")
            if missing["stockfish"]:
                print("      Legal moves missed: ")
                for move in missing["stockfish"]:
                    print(f"       - {move}: {positions['stockfish'][move]}")
            if missing["nimfish"]:
                print("      Illegal moves generated: ")
                for move in missing["nimfish"]:
                    print(f"       - {move}: {positions['nimfish'][move]}")
        if mistakes:
            print("   Counting mistakes made:")
            for move in mistakes:
                missed = positions["stockfish"][move] - positions["nimfish"][move]
                print(f"       - {move}: expected {positions['stockfish'][move]}, got {positions['nimfish'][move]} ({'-' if missed > 0 else '+'}{abs(missed)})")
    else:
        print("No discrepancies detected")

    


if __name__ == "__main__":
    parser = ArgumentParser(description="Automatically compare perft results between our engine and Stockfish")
    parser.add_argument("--fen", "-f", type=str, default="", help="The FEN string of the position to start from (empty string means the initial one). Defaults to ''")
    parser.add_argument("--ply", "-d", type=int, required=True, help="The depth to stop at, expressed in plys (half-moves)")
    parser.add_argument("--bulk", action="store_true", help="Enable bulk-counting for Nimfish (faster, less debuggable)", default=False)
    parser.add_argument("--stockfish", type=Path, help="Path to the stockfish executable. Defaults to '' (detected automatically)", default=None)
    parser.add_argument("--nimfish", type=Path, help="Path to the nimfish executable. Defaults to 'nimfish'", default=Path("nimfish"))
    sys.exit(main(parser.parse_args()))