optirs-learned 0.3.1

// Reinforcement learning-based neural architecture search
//
// This module implements RL agents that learn to design optimal optimizer architectures
// through interaction with the optimization environment and reward signals.

#[allow(dead_code)]

use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::numeric::Float;
use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;

use crate::error::{OptimError, Result};

/// Reinforcement learning configuration for NAS
#[derive(Debug, Clone)]
pub struct RLNASConfig<T: Float + Debug + Send + Sync + 'static> {
    /// Learning rate for policy updates
    pub learning_rate: T,
    
    /// Discount factor for future rewards
    pub discount_factor: T,
    
    /// Exploration rate (epsilon)
    pub exploration_rate: T,
    
    /// Exploration decay rate
    pub exploration_decay: T,
    
    /// Minimum exploration rate
    pub min_exploration_rate: T,
    
    /// Replay buffer size
    pub replay_buffer_size: usize,
    
    /// Batch size for training
    pub batch_size: usize,
    
    /// Number of training episodes
    pub num_episodes: usize,
    
    /// Maximum steps per episode
    pub max_steps_per_episode: usize,
    
    /// Target network update frequency
    pub target_update_frequency: usize,
    
    /// Architecture complexity penalty
    pub complexity_penalty: T,
}

/// Reinforcement learning agent for architecture search
#[derive(Debug)]
pub struct RLArchitectureAgent<T: Float + Debug + Send + Sync + 'static> {
    /// Agent configuration
    config: RLNASConfig<T>,
    
    /// Policy network for action selection
    policy_network: PolicyNetwork<T>,
    
    /// Value network for value estimation
    value_network: ValueNetwork<T>,
    
    /// Target networks for stable training
    target_policy_network: Option<PolicyNetwork<T>>,
    target_value_network: Option<ValueNetwork<T>>,
    
    /// Experience replay buffer
    replay_buffer: ReplayBuffer<T>,
    
    /// Current state representation
    current_state: Option<StateRepresentation<T>>,
    
    /// Training statistics
    training_stats: TrainingStatistics<T>,
    
    /// Episode history
    episode_history: Vec<Episode<T>>,
    
    /// Current episode
    current_episode: usize,
}

/// Policy network for action selection
#[derive(Debug, Clone)]
pub struct PolicyNetwork<T: Float + Debug + Send + Sync + 'static> {
    /// Network weights
    weights: Vec<Array2<T>>,
    
    /// Biases
    biases: Vec<Array1<T>>,
    
    /// Network architecture
    architecture: Vec<usize>,
    
    /// Output action space
    action_space: ActionSpace,
}

/// Value network for state value estimation
#[derive(Debug, Clone)]
pub struct ValueNetwork<T: Float + Debug + Send + Sync + 'static> {
    /// Network weights
    weights: Vec<Array2<T>>,
    
    /// Biases
    biases: Vec<Array1<T>>,
    
    /// Network architecture
    architecture: Vec<usize>,
}

/// Action space for architecture construction
#[derive(Debug, Clone)]
pub struct ActionSpace {
    /// Available layer types
    layer_types: Vec<LayerTypeAction>,
    
    /// Available connections
    connection_types: Vec<ConnectionAction>,
    
    /// Parameter ranges
    parameter_ranges: HashMap<String, (f64, f64)>,
    
    /// Maximum architecture depth
    max_depth: usize,
}

/// Layer type actions
#[derive(Debug, Clone, Copy)]
pub enum LayerTypeAction {
    AddDense { units: usize },
    AddConv1D { filters: usize, kernel_size: usize },
    AddLSTM { units: usize },
    AddAttention { heads: usize, dim: usize },
    AddBatchNorm,
    AddDropout { rate: f32 },
    AddActivation { activation_type: ActivationType },
    TerminateArchitecture,
}

/// Connection actions
#[derive(Debug, Clone, Copy)]
pub enum ConnectionAction {
    SequentialConnection,
    SkipConnection { from_layer: usize, to_layer: usize },
    DenseConnection,
    NoConnection,
}

/// Activation types for layers
#[derive(Debug, Clone, Copy)]
pub enum ActivationType {
    ReLU,
    Tanh,
    Sigmoid,
    GELU,
    Swish,
}

/// State representation of current architecture being built
#[derive(Debug, Clone)]
pub struct StateRepresentation<T: Float + Debug + Send + Sync + 'static> {
    /// Current layers in architecture
    current_layers: Vec<LayerInfo>,
    
    /// Current connections
    current_connections: Vec<(usize, usize)>,
    
    /// Architecture complexity metrics
    complexity_metrics: ComplexityMetrics<T>,
    
    /// Performance history
    performance_history: Vec<T>,
    
    /// State embedding
    state_embedding: Array1<T>,
}

/// Layer information in state
#[derive(Debug, Clone)]
pub struct LayerInfo {
    /// Layer type
    layer_type: String,
    
    /// Input dimensions
    input_dims: Vec<usize>,
    
    /// Output dimensions
    output_dims: Vec<usize>,
    
    /// Parameters
    parameters: HashMap<String, f64>,
}

/// Architecture complexity metrics
#[derive(Debug, Clone)]
pub struct ComplexityMetrics<T: Float + Debug + Send + Sync + 'static> {
    /// Number of parameters
    parameter_count: usize,
    
    /// Estimated FLOPs
    flops: u64,
    
    /// Memory requirements
    memory_usage: usize,
    
    /// Architectural depth
    depth: usize,
    
    /// Connectivity complexity
    connectivity_complexity: T,
}

/// Experience replay buffer
#[derive(Debug)]
pub struct ReplayBuffer<T: Float + Debug + Send + Sync + 'static> {
    /// Buffer capacity
    capacity: usize,
    
    /// Stored experiences
    experiences: VecDeque<Experience<T>>,
    
    /// Current position in buffer
    position: usize,
}

/// Experience tuple for replay
#[derive(Debug, Clone)]
pub struct Experience<T: Float + Debug + Send + Sync + 'static> {
    /// State before action
    state: StateRepresentation<T>,
    
    /// Action taken
    action: Action,
    
    /// Reward received
    reward: T,
    
    /// Next state
    next_state: Option<StateRepresentation<T>>,
    
    /// Whether episode terminated
    done: bool,
}

/// Action taken by agent
#[derive(Debug, Clone)]
pub enum Action {
    /// Layer type action
    LayerAction(LayerTypeAction),
    
    /// Connection action
    ConnectionAction(ConnectionAction),
    
    /// Hyperparameter adjustment
    ParameterAction { parameter: String, value: f64 },
}

/// Training statistics for RL agent
#[derive(Debug, Clone)]
pub struct TrainingStatistics<T: Float + Debug + Send + Sync + 'static> {
    /// Episode rewards
    episode_rewards: Vec<T>,
    
    /// Average rewards over time
    average_rewards: Vec<T>,
    
    /// Policy losses
    policy_losses: Vec<T>,
    
    /// Value losses
    value_losses: Vec<T>,
    
    /// Exploration rates over time
    exploration_rates: Vec<T>,
    
    /// Architecture performance history
    architecture_performances: Vec<T>,
}

/// Episode information
#[derive(Debug, Clone)]
pub struct Episode<T: Float + Debug + Send + Sync + 'static> {
    /// Episode number
    episode_number: usize,
    
    /// Steps taken
    steps: usize,
    
    /// Total reward
    total_reward: T,
    
    /// Final architecture
    final_architecture: Option<StateRepresentation<T>>,
    
    /// Performance achieved
    performance: Option<T>,
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> RLArchitectureAgent<T> {
    /// Create new RL architecture agent
    pub fn new(config: RLNASConfig<T>) -> Result<Self> {
        let policy_network = PolicyNetwork::new(vec![128, 256, 256, 64])?;
        let value_network = ValueNetwork::new(vec![128, 256, 128, 1])?;
        let replay_buffer = ReplayBuffer::new(config.replay_buffer_size);
        let training_stats = TrainingStatistics::new();
        
        Ok(Self {
            config,
            policy_network,
            value_network,
            target_policy_network: None,
            target_value_network: None,
            replay_buffer,
            current_state: None,
            training_stats,
            episode_history: Vec::new(),
            current_episode: 0,
        })
    }

    /// Train the RL agent to learn architecture design
    pub fn train(&mut self, env: &mut dyn ArchitectureEnvironment<T>) -> Result<()> {
        // Initialize target networks
        self.target_policy_network = Some(self.policy_network.clone());
        self.target_value_network = Some(self.value_network.clone());
        
        for episode in 0..self.config.num_episodes {
            self.current_episode = episode;
            
            // Reset environment and get initial state
            let initial_state = env.reset()?;
            self.current_state = Some(initial_state);
            
            let mut episode_reward = T::zero();
            let mut steps = 0;
            
            // Run episode
            for step in 0..self.config.max_steps_per_episode {
                steps = step + 1;
                
                // Select action using epsilon-greedy policy
                let action = self.select_action(self.current_state.as_ref().expect("unwrap failed"))?;
                
                // Execute action in environment
                let (next_state, reward, done) = env.step(&action)?;
                episode_reward = episode_reward + reward;
                
                // Store experience in replay buffer
                let experience = Experience {
                    state: self.current_state.as_ref().expect("unwrap failed").clone(),
                    action,
                    reward,
                    next_state: if done { None } else { Some(next_state.clone()) },
                    done,
                };
                self.replay_buffer.add(experience);
                
                // Update current state
                self.current_state = if done { None } else { Some(next_state) };
                
                // Train networks if enough experiences collected
                if self.replay_buffer.size() >= self.config.batch_size {
                    self.train_networks()?;
                }
                
                if done {
                    break;
                }
            }
            
            // Record episode statistics
            self.training_stats.episode_rewards.push(episode_reward);
            self.update_training_statistics();
            
            // Update target networks periodically
            if episode % self.config.target_update_frequency == 0 {
                self.update_target_networks();
            }
            
            // Decay exploration rate
            self.decay_exploration_rate();
            
            // Log progress
            if episode % 100 == 0 {
                let avg_reward = self.training_stats.episode_rewards.iter()
                    .rev()
                    .take(100)
                    .cloned()
                    .fold(T::zero(), |acc, r| acc + r) / T::from(100.0.min(episode as f64 + 1.0)).expect("unwrap failed");
                println!("Episode {}: Average reward = {:.4}", episode, avg_reward.to_f64().unwrap_or(0.0));
            }
        }
        
        Ok(())
    }

    /// Select action using epsilon-greedy policy
    fn select_action(&self, state: &StateRepresentation<T>) -> Result<Action> {
        use scirs2_core::random::Rng;
        let mut rng = scirs2_core::random::thread_rng();
        
        // Epsilon-greedy exploration
        if rng.random::<f64>() < self.config.exploration_rate.to_f64().unwrap_or(0.1) {
            // Random action
            self.select_random_action(state)
        } else {
            // Policy-based action
            self.select_policy_action(state)
        }
    }

    /// Select random action for exploration
    fn select_random_action(&self, state: &StateRepresentation<T>) -> Result<Action> {
        use scirs2_core::random::Rng;
        let mut rng = scirs2_core::random::thread_rng();
        
        // Randomly choose action type
        let action_type = rng.gen_range(0..3);
        
        match action_type {
            0 => {
                // Layer action
                let layer_types = vec![
                    LayerTypeAction::AddDense { units: rng.gen_range(32..512) },
                    LayerTypeAction::AddLSTM { units: rng.gen_range(64..256) },
                    LayerTypeAction::AddAttention { heads: rng.gen_range(4..16)..dim: rng.gen_range(64..256) },
                    LayerTypeAction::AddBatchNorm,
                    LayerTypeAction::AddDropout { rate: rng.gen_range(0.1..0.5) },
                ];
                let layer_action = layer_types[rng.gen_range(0..layer_types.len())];
                Ok(Action::LayerAction(layer_action))
            }
            1 => {
                // Connection action
                let conn_action = if state.current_layers.len() > 1 {
                    ConnectionAction::SkipConnection { 
                        from_layer: rng.gen_range(0..state.current_layers.len()-1),
                        to_layer: rng.gen_range(1..state.current_layers.len())
                    }
                } else {
                    ConnectionAction::SequentialConnection
                };
                Ok(Action::ConnectionAction(conn_action))
            }
            _ => {
                // Parameter action
                Ok(Action::ParameterAction { 
                    parameter: "learning_rate".to_string(),
                    value: rng.gen_range(0.0001..0.01)
                })
            }
        }
    }

    /// Select action using policy network
    fn select_policy_action(&self, state: &StateRepresentation<T>) -> Result<Action> {
        // Forward pass through policy network
        let policy_output = self.policy_network.forward(&state.state_embedding)?;
        
        // Select action with highest probability
        let (max_idx, _max_val) = policy_output.iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .unwrap_or((0, &T::zero()));
        
        // Map index to action (simplified mapping)
        match max_idx % 8 {
            0 => Ok(Action::LayerAction(LayerTypeAction::AddDense { units: 128 })),
            1 => Ok(Action::LayerAction(LayerTypeAction::AddLSTM { units: 128 })),
            2 => Ok(Action::LayerAction(LayerTypeAction::AddAttention { heads: 8, dim: 128 })),
            3 => Ok(Action::LayerAction(LayerTypeAction::AddBatchNorm)),
            4 => Ok(Action::LayerAction(LayerTypeAction::AddDropout { rate: 0.2 })),
            5 => Ok(Action::ConnectionAction(ConnectionAction::SequentialConnection)),
            6 => Ok(Action::ConnectionAction(ConnectionAction::DenseConnection)),
            _ => Ok(Action::LayerAction(LayerTypeAction::TerminateArchitecture)),
        }
    }

    /// Train policy and value networks
    fn train_networks(&mut self) -> Result<()> {
        // Sample batch from replay buffer
        let batch = self.replay_buffer.sample(self.config.batch_size)?;
        
        // Compute policy loss and update
        let policy_loss = self.compute_policy_loss(&batch)?;
        self.update_policy_network(policy_loss)?;
        
        // Compute value loss and update
        let value_loss = self.compute_value_loss(&batch)?;
        self.update_value_network(value_loss)?;
        
        // Record losses
        self.training_stats.policy_losses.push(policy_loss);
        self.training_stats.value_losses.push(value_loss);
        
        Ok(())
    }

    /// Compute policy loss (simplified)
    fn compute_policy_loss(&self, _batch: &[Experience<T>]) -> Result<T> {
        // Simplified policy loss computation
        // In practice, this would involve computing advantage estimates
        // and policy gradient loss
        Ok(scirs2_core::numeric::NumCast::from(0.1).unwrap_or_else(|| T::zero()))
    }

    /// Compute value loss (simplified)
    fn compute_value_loss(&self, _batch: &[Experience<T>]) -> Result<T> {
        // Simplified value loss computation
        // In practice, this would compute temporal difference error
        Ok(scirs2_core::numeric::NumCast::from(0.05).unwrap_or_else(|| T::zero()))
    }

    /// Update policy network (simplified)
    fn update_policy_network(&mut self, _loss: T) -> Result<()> {
        // Simplified network update
        // In practice, this would apply gradient descent
        Ok(())
    }

    /// Update value network (simplified)
    fn update_value_network(&mut self, _loss: T) -> Result<()> {
        // Simplified network update
        // In practice, this would apply gradient descent
        Ok(())
    }

    /// Update target networks
    fn update_target_networks(&mut self) {
        self.target_policy_network = Some(self.policy_network.clone());
        self.target_value_network = Some(self.value_network.clone());
    }

    /// Decay exploration rate
    fn decay_exploration_rate(&mut self) {
        let current_rate = self.config.exploration_rate;
        let decayed_rate = current_rate * self.config.exploration_decay;
        self.config.exploration_rate = decayed_rate.max(self.config.min_exploration_rate);
        self.training_stats.exploration_rates.push(self.config.exploration_rate);
    }

    /// Update training statistics
    fn update_training_statistics(&mut self) {
        if !self.training_stats.episode_rewards.is_empty() {
            let window_size = 100.min(self.training_stats.episode_rewards.len());
            let recent_rewards = &self.training_stats.episode_rewards[
                self.training_stats.episode_rewards.len().saturating_sub(window_size)..
            ];
            
            let average_reward = recent_rewards.iter().cloned().fold(T::zero(), |acc, r| acc + r) /
                T::from(recent_rewards.len() as f64).expect("unwrap failed");
            
            self.training_stats.average_rewards.push(average_reward);
        }
    }

    /// Get training statistics
    pub fn get_training_statistics(&self) -> &TrainingStatistics<T> {
        &self.training_stats
    }

    /// Get best architecture found during training
    pub fn get_best_architecture(&self) -> Option<StateRepresentation<T>> {
        self.episode_history.iter()
            .max_by(|a, b| a.performance.unwrap_or(T::zero()).partial_cmp(&b.performance.unwrap_or(T::zero())).unwrap_or(std::cmp::Ordering::Equal))
            .and_then(|episode| episode.final_architecture.clone())
    }
}

/// Environment interface for RL training
pub trait ArchitectureEnvironment<T: Float + Debug + Send + Sync + 'static> {
    /// Reset environment and return initial state
    fn reset(&mut self) -> Result<StateRepresentation<T>>;
    
    /// Execute action and return (next_state, reward, done)
    fn step(&mut self, action: &Action) -> Result<(StateRepresentation<T>, T, bool)>;
    
    /// Get current state
    fn get_state(&self) -> Result<StateRepresentation<T>>;
    
    /// Evaluate current architecture
    fn evaluate_architecture(&self, state: &StateRepresentation<T>) -> Result<T>;
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> PolicyNetwork<T> {
    fn new(architecture: Vec<usize>) -> Result<Self> {
        let mut weights = Vec::new();
        let mut biases = Vec::new();
        
        for i in 0..architecture.len()-1 {
            let w = Array2::zeros((architecture[i], architecture[i+1]));
            let b = Array1::zeros(architecture[i+1]);
            weights.push(w);
            biases.push(b);
        }
        
        Ok(Self {
            weights,
            biases,
            architecture: architecture.clone(),
            action_space: ActionSpace::default(),
        })
    }
    
    fn forward(&self, input: &Array1<T>) -> Result<Array1<T>> {
        let mut activation = input.clone();
        
        for (w, b) in self.weights.iter().zip(self.biases.iter()) {
            let mut output = Array1::zeros(w.ncols());
            for i in 0..w.ncols() {
                for j in 0..w.nrows() {
                    output[i] = output[i] + activation[j] * w[[j, i]];
                }
                output[i] = output[i] + b[i];
            }
            
            // Apply activation function (ReLU)
            activation = output.mapv(|x| x.max(T::zero()));
        }
        
        Ok(activation)
    }
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> ValueNetwork<T> {
    fn new(architecture: Vec<usize>) -> Result<Self> {
        let mut weights = Vec::new();
        let mut biases = Vec::new();
        
        for i in 0..architecture.len()-1 {
            let w = Array2::zeros((architecture[i], architecture[i+1]));
            let b = Array1::zeros(architecture[i+1]);
            weights.push(w);
            biases.push(b);
        }
        
        Ok(Self {
            weights,
            biases,
            architecture,
        })
    }
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> ReplayBuffer<T> {
    fn new(capacity: usize) -> Self {
        Self {
            capacity,
            experiences: VecDeque::new(),
            position: 0,
        }
    }
    
    fn add(&mut self, experience: Experience<T>) {
        if self.experiences.len() < self.capacity {
            self.experiences.push_back(experience);
        } else {
            self.experiences[self.position] = experience;
            self.position = (self.position + 1) % self.capacity;
        }
    }
    
    fn sample(&self, batch_size: usize) -> Result<Vec<Experience<T>>> {
        use scirs2_core::random::Rng;
        if self.experiences.len() < batch_size {
            return Err(OptimError::InsufficientData(
                "Not enough experiences in replay buffer".to_string()
            ));
        }
        
        let mut rng = scirs2_core::random::thread_rng();
        let mut batch = Vec::with_capacity(batch_size);
        
        for _ in 0..batch_size {
            let idx = rng.gen_range(0..self.experiences.len());
            batch.push(self.experiences[idx].clone());
        }
        
        Ok(batch)
    }
    
    fn size(&self) -> usize {
        self.experiences.len()
    }
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> TrainingStatistics<T> {
    fn new() -> Self {
        Self {
            episode_rewards: Vec::new(),
            average_rewards: Vec::new(),
            policy_losses: Vec::new(),
            value_losses: Vec::new(),
            exploration_rates: Vec::new(),
            architecture_performances: Vec::new(),
        }
    }
}

impl Default for ActionSpace {
    fn default() -> Self {
        let mut parameter_ranges = HashMap::new();
        parameter_ranges.insert("units".to_string(), (32.0, 512.0));
        parameter_ranges.insert("filters".to_string(), (16.0, 256.0));
        parameter_ranges.insert("kernel_size".to_string(), (1.0, 7.0));
        parameter_ranges.insert("dropout_rate".to_string(), (0.1, 0.5));
        
        Self {
            layer_types: vec![
                LayerTypeAction::AddDense { units: 128 },
                LayerTypeAction::AddLSTM { units: 128 },
                LayerTypeAction::AddAttention { heads: 8, dim: 128 },
            ],
            connection_types: vec![
                ConnectionAction::SequentialConnection,
                ConnectionAction::DenseConnection,
            ],
            parameter_ranges,
            max_depth: 20,
        }
    }
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> Default for RLNASConfig<T> {
    fn default() -> Self {
        Self {
            learning_rate: scirs2_core::numeric::NumCast::from(0.001).unwrap_or_else(|| T::zero()),
            discount_factor: scirs2_core::numeric::NumCast::from(0.99).unwrap_or_else(|| T::zero()),
            exploration_rate: scirs2_core::numeric::NumCast::from(1.0).unwrap_or_else(|| T::zero()),
            exploration_decay: scirs2_core::numeric::NumCast::from(0.995).unwrap_or_else(|| T::zero()),
            min_exploration_rate: scirs2_core::numeric::NumCast::from(0.01).unwrap_or_else(|| T::zero()),
            replay_buffer_size: 10000,
            batch_size: 32,
            num_episodes: 1000,
            max_steps_per_episode: 50,
            target_update_frequency: 100,
            complexity_penalty: scirs2_core::numeric::NumCast::from(0.01).unwrap_or_else(|| T::zero()),
        }
    }
}