#!/usr/bin/env python3
"""Generate random chess positions for NNUE training with minimal filtering."""

import chess
import random
import sys
from pathlib import Path
from tqdm import tqdm

def play_random_game_and_collect_positions(
    output_file,
    total_games=500000,
    samples_per_game=1,
    min_move=1,
    max_move=50
):
    """Play random games and sample multiple positions from each.

    Args:
        output_file: Output file for positions
        total_games: Number of games to play
        samples_per_game: Number of positions to sample per game (1-N)
        min_move: Minimum move number to start sampling from
        max_move: Maximum move number for sampling

    Returns:
        Number of valid positions saved
    """
    positions_count = 0
    filtered_game_over = 0
    filtered_illegal = 0

    with open(output_file, 'w') as f:
        with tqdm(total=total_games, desc="Generating positions") as pbar:
            for game_num in range(total_games):
                board = chess.Board()
                move_history = []

                # Play a complete random game
                while not board.is_game_over() and len(move_history) < 200:
                    legal_moves = list(board.legal_moves)
                    if not legal_moves:
                        break
                    move = random.choice(legal_moves)
                    board.push(move)
                    move_history.append(board.copy())

                # Determine the range of moves to sample from
                game_length = len(move_history)
                valid_start = max(min_move, 0)
                valid_end = min(max_move, game_length)

                if valid_start >= valid_end:
                    # Game too short
                    pbar.update(1)
                    continue

                # Randomly sample positions from this game
                sample_count = min(samples_per_game, valid_end - valid_start)
                if sample_count > 0:
                    sample_indices = random.sample(
                        range(valid_start, valid_end),
                        k=sample_count
                    )

                    for idx in sample_indices:
                        sampled_board = move_history[idx]

                        # Only filter truly invalid or terminal positions
                        if not sampled_board.is_valid():
                            filtered_illegal += 1
                            continue

                        if sampled_board.is_game_over():
                            filtered_game_over += 1
                            continue

                        # Save position (include check, captures, all positions)
                        fen = sampled_board.fen()
                        f.write(fen + '\n')
                        positions_count += 1

                pbar.update(1)

    # Print summary
    print()
    print("=" * 60)
    print("POSITION GENERATION SUMMARY")
    print("=" * 60)
    print(f"Total games played:        {total_games}")
    print(f"Samples per game:          {samples_per_game}")
    print(f"Move range:                {min_move}-{max_move}")
    print(f"Saved positions:           {positions_count}")
    print(f"Filtered (game over):      {filtered_game_over}")
    print(f"Filtered (illegal):        {filtered_illegal}")
    print(f"Acceptance rate:           {positions_count / total_games * 100:.2f}%")
    print(f"(Includes checks, captures, all realistic positions)")
    print("=" * 60)
    print()

    if positions_count == 0:
        print("WARNING: No valid positions were generated!")
        return 0

    return positions_count

if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Generate random chess positions for NNUE training")
    parser.add_argument("output_file", nargs="?", default="positions.txt",
                        help="Output file for positions (default: positions.txt)")
    parser.add_argument("--games", type=int, default=5000,
                        help="Number of games to play (default: 5000)")
    parser.add_argument("--samples-per-game", type=int, default=1,
                        help="Number of positions to sample per game (default: 1)")
    parser.add_argument("--min-move", type=int, default=1,
                        help="Minimum move number to sample from (default: 1)")
    parser.add_argument("--max-move", type=int, default=50,
                        help="Maximum move number to sample from (default: 50)")

    args = parser.parse_args()

    count = play_random_game_and_collect_positions(
        output_file=args.output_file,
        total_games=args.games,
        samples_per_game=args.samples_per_game,
        min_move=args.min_move,
        max_move=args.max_move
    )

    sys.exit(0 if count > 0 else 1)
