#!/usr/bin/env python3
"""Generate 500,000 random chess positions for NNUE training."""

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

# Standard piece values for capture filtering
PIECE_VALUES = {
    chess.PAWN: 1,
    chess.KNIGHT: 3,
    chess.BISHOP: 3,
    chess.ROOK: 5,
    chess.QUEEN: 9,
}

def has_winning_or_equal_capture(board):
    """Check if position has a capture where victim >= attacker (winning or equal trade).

    Returns True only if there's at least one favorable capture.
    Positions with only losing captures return False (are kept).
    """
    for move in board.legal_moves:
        if board.is_capture(move):
            attacker_piece = board.piece_at(move.from_square)
            victim_piece = board.piece_at(move.to_square)

            if attacker_piece and victim_piece:
                attacker_value = PIECE_VALUES.get(attacker_piece.piece_type, 0)
                victim_value = PIECE_VALUES.get(victim_piece.piece_type, 0)

                # If victim >= attacker, it's a winning or equal capture
                if victim_value >= attacker_value:
                    return True

    return False

def play_random_game_and_collect_positions(output_file, total_games=500000, filter_captures=True):
    """Play random games and save positions after 8-20 random moves.

    Returns:
        Number of valid positions saved
    """
    positions_count = 0
    filtered_check = 0
    filtered_captures = 0
    filtered_game_over = 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()

                # Play 8-20 random opening moves
                num_moves = random.randint(8, 20)

                for move_num in range(num_moves):
                    if board.is_game_over():
                        break

                    legal_moves = list(board.legal_moves)
                    if not legal_moves:
                        break

                    move = random.choice(legal_moves)
                    board.push(move)

                # Skip if game over
                if board.is_game_over():
                    filtered_game_over += 1
                    pbar.update(1)
                    continue

                # Skip if in check
                if board.is_check():
                    filtered_check += 1
                    pbar.update(1)
                    continue

                # Check if there are winning or equal captures (if filtering enabled)
                if filter_captures:
                    if has_winning_or_equal_capture(board):
                        filtered_captures += 1
                        pbar.update(1)
                        continue

                # Save valid position
                fen = board.fen()
                f.write(fen + '\n')
                positions_count += 1

                pbar.update(1)

    # Print summary
    print()
    print("=" * 60)
    print("POSITION GENERATION SUMMARY")
    print("=" * 60)
    total_filtered = filtered_check + filtered_captures + filtered_game_over
    print(f"Total games:               {total_games}")
    print(f"Saved positions:           {positions_count}")
    print(f"Filtered (in check):       {filtered_check}")
    print(f"Filtered (winning+ cap):   {filtered_captures}")
    print(f"Filtered (game over):      {filtered_game_over}")
    print(f"Total filtered:            {total_filtered}")
    print(f"Acceptance rate:           {positions_count / total_games * 100:.2f}%")
    print(f"(Keeps positions with only losing/bad captures)")
    print("=" * 60)
    print()

    if positions_count == 0:
        print("WARNING: No valid positions were generated!")
        print("This might happen if:")
        print("  - Most positions have checks or game-over states")
        print("  - Try using: --no-filter-captures to accept positions with winning captures")
        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: 500000)")
    parser.add_argument("--no-filter-captures", action="store_true",
                        help="Include positions with winning/equal captures (increases output)")

    args = parser.parse_args()

    count = play_random_game_and_collect_positions(
        output_file=args.output_file,
        total_games=args.games,
        filter_captures=not args.no_filter_captures
    )

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