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"Saved {saved} positions.")


if __name__ == "__main__":
    main()