1 files changed, 55 insertions, 0 deletions
diff --git a/match-manager/extract.py b/match-manager/extract.py
new file mode 100644
index 0000000..49c57a6
--- /dev/null
+++ b/match-manager/extract.py
@@ -0,0 +1,55 @@
+import random
+
+import chess.pgn
+
+INPUT = "big_games.pgn"
+OUTPUT = "positions.txt"
+
+
+def get_random_position(game):
+    moves = list(game.mainline_moves())
+
+    # Game is too short
+    if len(moves) <= 10:
+        return None
+
+    move_number = random.randint(10, min(len(moves) - 1, 40))
+
+    board = game.board()
+
+    for i, move in enumerate(moves):
+        board.push(move)
+        if i == move_number:
+            break
+
+    return board.fen()
+
+
+def main():
+    max_games = int(input("Enter number of game to extract: "))
+
+    saved = 0
+
+    with open(INPUT, "r") as pgn_file, open(OUTPUT, "w") as fen_file:
+        for i in range(max_games):
+            game = chess.pgn.read_game(pgn_file)
+
+            if game is None:
+                break
+
+            fen = get_random_position(game)
+
+            if fen is None:
+                continue
+
+            fen_file.write(fen + "\n")
+
+            saved += 1
+
+            print(f"{i + 1}/{max_games} > saved {saved}")
+
+    print(f"Done. Saved {saved} positions.")
+
+
+if __name__ == "__main__":
+    main()