"""AlphaZero configuration."""

from __future__ import annotations

from dataclasses import dataclass

import torch


@dataclass
class AlphaZeroConfig:
    """Configuration for AlphaZero training.

    Attributes:
        num_simulations: MCTS simulations per move
        c_puct: Exploration constant for PUCT formula
        dirichlet_alpha: Dirichlet noise parameter for root exploration
        dirichlet_epsilon: Weight of Dirichlet noise (1-epsilon is prior weight)
        mcts_batch_size: Batch size for GPU inference during MCTS
        mcts_virtual_loss: Virtual loss for batched MCTS
        num_iterations: Number of training iterations
        games_per_iteration: Self-play games per iteration
        training_steps_per_iteration: Gradient steps per iteration
        batch_size: Training batch size
        learning_rate: Optimizer learning rate
        weight_decay: L2 regularization weight
        replay_buffer_size: Maximum replay buffer capacity
        min_replay_size: Minimum samples before training starts
        bc_ratio: Ratio of BC samples in training batches (for dual buffer mode)
        hidden_dim: Network hidden dimension
        num_blocks: Number of residual blocks in network
        temperature_moves: Use temperature=1 for first N moves
        temperature_final: Temperature after N moves
        eval_interval: Evaluate every N iterations
        eval_games: Number of games for evaluation
        checkpoint_interval: Save checkpoint every N iterations (0 = only at end)
        device: Device for training (cuda/cpu)
        normalize_obs: Whether to normalize observations
    """

    # MCTS
    num_simulations: int = 100
    c_puct: float = 1.5
    dirichlet_alpha: float = 0.3
    dirichlet_epsilon: float = 0.25
    mcts_batch_size: int = 32
    mcts_virtual_loss: float = 3.0

    # Training
    num_iterations: int = 100
    games_per_iteration: int = 100
    training_steps_per_iteration: int = 100
    batch_size: int = 256
    learning_rate: float = 1e-3
    weight_decay: float = 1e-4

    # Replay buffer
    replay_buffer_size: int = 100_000
    min_replay_size: int = 1000

    # BC warm-start (dual buffer mode)
    bc_ratio: float = 0.7

    # Network
    hidden_dim: int = 256
    num_blocks: int = 4

    # Temperature for action selection
    temperature_moves: int = 30
    temperature_final: float = 0.1

    # Evaluation
    eval_interval: int = 5
    eval_games: int = 50

    # Checkpointing
    checkpoint_interval: int = 10

    # Device
    device: str = "cuda" if torch.cuda.is_available() else "cpu"

    # Observation normalization
    normalize_obs: bool = True
