//! Rainbow DQN Configuration Types
//!
//! Configuration structures for Rainbow DQN agent and its components

use serde::{Deserialize, Serialize};

use super::distributional::DistributionalConfig;
use super::rainbow_network::RainbowNetworkConfig;

/// Configuration for Rainbow `DQN` Agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RainbowAgentConfig {
    /// Device to run on ("cpu" or "cuda")
    pub device: String,
    /// Network configuration
    pub network_config: RainbowNetworkConfig,
    /// Minimum replay buffer size before training
    pub min_replay_size: usize,
    /// Experience replay buffer capacity
    pub replay_buffer_size: usize,
    /// Training batch size
    pub batch_size: usize,
    /// Learning rate
    pub learning_rate: f64,
    /// Discount factor
    pub gamma: f64,
    /// Target network update frequency
    pub target_update_freq: usize,
    /// Training frequency (steps between training)
    pub train_freq: usize,
    /// Priority replay configuration
    pub priority_alpha: f64,
    pub priority_beta: f64,
    pub priority_beta_increment: f64,
    /// Noisy network reset frequency
    pub noise_reset_freq: usize,
}

impl Default for RainbowAgentConfig {
    fn default() -> Self {
        Self {
            device: "cuda".to_owned(),
            network_config: RainbowNetworkConfig::default(),
            min_replay_size: 10000,
            replay_buffer_size: 100000,
            batch_size: 32,
            learning_rate: 0.0001,
            gamma: 0.99,
            target_update_freq: 1000,
            train_freq: 4,
            priority_alpha: 0.6,
            priority_beta: 0.4,
            priority_beta_increment: 0.00025,
            noise_reset_freq: 100,
        }
    }
}

/// Metrics for Rainbow `DQN` Agent
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RainbowAgentMetrics {
    /// Total training steps
    pub total_steps: u64,
    /// Total episodes
    pub total_episodes: u64,
    /// Current replay buffer size
    pub replay_buffer_size: usize,
    /// Average reward over last 100 episodes
    pub average_reward: f64,
    /// Current loss value
    pub current_loss: f64,
    /// Current exploration rate
    pub exploration_rate: f64,
    /// Priority replay metrics
    pub priority_beta: f64,
    /// Training throughput (steps/second)
    pub training_throughput: f64,
}

impl Default for RainbowAgentMetrics {
    fn default() -> Self {
        Self {
            total_steps: 0,
            total_episodes: 0,
            replay_buffer_size: 0,
            average_reward: 0.0,
            current_loss: 0.0,
            exploration_rate: 1.0,
            priority_beta: 0.4,
            training_throughput: 0.0,
        }
    }
}

/// Configuration for Rainbow `DQN` system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RainbowDQNConfig {
    /// Network architecture configuration
    pub network: RainbowNetworkConfig,
    /// Distributional RL configuration
    pub distributional: DistributionalConfig,
    /// Training configuration
    pub learning_rate: f64,
    pub batch_size: usize,
    pub replay_buffer_size: usize,
    pub target_update_freq: usize,
    /// Priority replay configuration
    pub priority_replay_enabled: bool,
    pub priority_alpha: f64,
    pub priority_beta: f64,
    /// Device configuration
    pub device: String,
}

impl Default for RainbowDQNConfig {
    fn default() -> Self {
        Self {
            network: RainbowNetworkConfig {
                input_size: 10,
                hidden_sizes: vec![64, 64],
                num_actions: 3,
                activation: super::rainbow_network::ActivationType::ReLU,
                dropout_rate: 0.1,
                distributional: DistributionalConfig::default(),
                noisy_sigma_init: 0.5,
                dueling: true,
                use_spectral_norm: false,
                spectral_norm_iterations: 1,
            },
            distributional: DistributionalConfig {
                num_atoms: 51,
                v_min: -25.0,  // DSR Q-values: rewards ±2 with gamma=0.92 → Q ≈ ±25
                v_max: 25.0,   // DSR Q-values: rewards ±2 with gamma=0.92 → Q ≈ ±25
            },
            learning_rate: 0.0001,
            batch_size: 32,
            replay_buffer_size: 100000,
            target_update_freq: 1000,
            priority_replay_enabled: true,
            priority_alpha: 0.6,
            priority_beta: 0.4,
            device: "cuda".to_owned(),
        }
    }
}

/// Metrics for Rainbow `DQN` system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RainbowMetrics {
    /// Total training steps
    pub total_steps: u64,
    /// Total episodes completed
    pub total_episodes: u64,
    /// Average reward over recent episodes
    pub average_reward: f64,
    /// Current exploration rate (for noisy networks, this is noise scale)
    pub exploration_rate: f64,
    /// Recent training loss
    pub training_loss: f64,
    /// Q-value statistics
    pub q_value_mean: f64,
    pub q_value_std: f64,
    /// Priority replay statistics
    pub priority_weight_mean: f64,
    pub priority_weight_max: f64,
    /// Network utilization
    pub network_updates: u64,
    pub target_network_updates: u64,
}

impl RainbowMetrics {
    pub const fn new() -> Self {
        Self {
            total_steps: 0,
            total_episodes: 0,
            average_reward: 0.0,
            exploration_rate: 1.0,
            training_loss: 0.0,
            q_value_mean: 0.0,
            q_value_std: 0.0,
            priority_weight_mean: 1.0,
            priority_weight_max: 1.0,
            network_updates: 0,
            target_network_updates: 0,
        }
    }
}

impl Default for RainbowMetrics {
    fn default() -> Self {
        Self::new()
    }
}

/// Training result for Rainbow `DQN`
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingResult {
    /// Loss value from this training step
    pub loss: f64,
    /// Q-value statistics
    pub q_values: Vec<f64>,
    /// Priority weights used (if priority replay enabled)
    pub priority_weights: Option<Vec<f64>>,
    /// Gradient norms for monitoring
    pub gradient_norm: f64,
    /// Network output distributions (for distributional RL)
    pub distributions: Option<Vec<Vec<f64>>>,
}

impl TrainingResult {
    pub const fn new(loss: f64) -> Self {
        Self {
            loss,
            q_values: Vec::new(),
            priority_weights: None,
            gradient_norm: 0.0,
            distributions: None,
        }
    }
}

impl Default for TrainingResult {
    fn default() -> Self {
        Self::new(0.0)
    }
}
