#!/usr/bin/env python3
"""Drops duplicate games from one or more self-play PGN files, keyed on
(starting FEN, move text) -- not on PGN headers (Date/Event/Result
alone don't prove two games are the same game) and not on FEN alone
(two games can legitimately share a random starting position and then
diverge, since self-play move choice is not deterministic across
engine generations/threads -- those are NOT duplicates and must be
kept).

Motivation: self-play starts from randomized opening positions (see
docs/NNUE.md); with enough games it becomes likely the same random
starting FEN gets drawn more than once. If the *same* position then
also produced the *same* game (e.g. a re-run over an unchanged engine
binary, or two shards drawing the same seed), that's a real duplicate
and should not be trained on twice -- it would silently inflate that
game's influence on both the WDL-K fit (fit_k_and_label.py) and the
training loss itself.

Dedup key is computed directly from each game's raw PGN text (the
[FEN "..."] header plus a comment-and-whitespace-stripped movetext
string), not by parsing through python-chess -- this project's own
split_pgn() (run_nnue_pipeline.py) found full-PGN python-chess parsing
too slow to use as a bulk pre-pass over large corpora; a raw-text key
is exact for this purpose (identical moves from an identical FEN is
exactly the definition of "the same game") and avoids that cost
entirely.

Usage:
    python3 dedup_pgn.py ../pgn/stockfish*.pgn --out-dir runs/dedup
Writes <out-dir>/<stem>.dedup.pgn for each input, in the same relative
order games appeared in their source file (cross-file: first file on
the command line wins ties). Prints a per-file and total duplicate
count; does not modify the inputs.
"""
import argparse
import hashlib
import re
import sys
from pathlib import Path

COMMENT_RE = re.compile(rb"\{[^}]*\}")
WHITESPACE_RE = re.compile(rb"\s+")
FEN_RE = re.compile(rb'\[FEN\s+"([^"]*)"\]')


def iter_raw_games(path):
    """Yields raw byte blocks for each game in a PGN file, split on
    lines starting "[Event " -- same boundary convention as
    run_nnue_pipeline.py's split_pgn()."""
    with open(path, "rb") as f:
        data = f.read()
    offsets = [m.start() for m in re.finditer(rb"^\[Event ", data, re.MULTILINE)]
    offsets.append(len(data))
    for i in range(len(offsets) - 1):
        block = data[offsets[i]:offsets[i + 1]]
        if block.strip():
            yield block


def game_key(raw_game):
    """(fen, movetext) dedup key. Movetext is everything after the
    header block (first blank line), with PGN brace-comments and all
    whitespace stripped -- comments carry score/depth/PV annotations
    that are irrelevant to "is this the same game" and would otherwise
    make two runs that scored an identical game differently (e.g. a
    deeper search depth) look like different games, which is wrong:
    same moves played is the actual definition of "the same game"
    here, regardless of what score commentary got attached to it."""
    fen_match = FEN_RE.search(raw_game)
    fen = fen_match.group(1) if fen_match else b""
    blank = raw_game.find(b"\n\n")
    movetext = raw_game[blank:] if blank != -1 else b""
    movetext = COMMENT_RE.sub(b"", movetext)
    movetext = WHITESPACE_RE.sub(b"", movetext)
    return hashlib.sha256(fen + b"|" + movetext).digest()


def main():
    ap = argparse.ArgumentParser(description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pgns", nargs="+", help="input PGN file(s)")
    ap.add_argument("--out-dir", required=True,
                     help="directory to write <stem>.dedup.pgn into")
    args = ap.parse_args()

    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    seen = set()
    total_in = 0
    total_dupes = 0
    for pgn_path in args.pgns:
        pgn_path = Path(pgn_path)
        n_in = 0
        n_dupes = 0
        out_path = out_dir / f"{pgn_path.stem}.dedup.pgn"
        with open(out_path, "wb") as out:
            for raw_game in iter_raw_games(pgn_path):
                n_in += 1
                key = game_key(raw_game)
                if key in seen:
                    n_dupes += 1
                    continue
                seen.add(key)
                out.write(raw_game)
                if not raw_game.endswith(b"\n\n"):
                    out.write(b"\n")
        print(f"{pgn_path}: {n_in} games, {n_dupes} duplicates dropped "
              f"-> {out_path} ({n_in - n_dupes} kept)")
        total_in += n_in
        total_dupes += n_dupes

    print(f"\nTotal: {total_in} games in, {total_dupes} duplicates "
          f"dropped, {total_in - total_dupes} kept across "
          f"{len(args.pgns)} file(s).")
    if total_dupes == 0:
        print("(no duplicates found)")


if __name__ == "__main__":
    sys.exit(main())