optirs-learned 0.3.1

use std::fmt::Debug;
// Reinforcement learning for neural architecture search
//
// This module implements RL-based NAS using controller networks,
// policy gradients, and other reinforcement learning methods.

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

use super::super::architecture::{ArchitectureSpec, ArchitectureCandidate, ActivationType};

/// Reinforcement learning search state
#[derive(Debug)]
pub struct RLSearchState<T: Float + Debug + Send + Sync + 'static> {
    /// Controller network
    pub controller: ControllerNetwork<T>,

    /// Action space
    pub action_space: ActionSpace,

    /// State representation
    pub state_representation: StateRepresentation<T>,

    /// Reward history
    pub reward_history: VecDeque<f64>,

    /// Policy parameters
    pub policy_parameters: PolicyParameters<T>,

    /// Experience replay buffer
    pub replay_buffer: ExperienceReplayBuffer<T>,

    /// Training configuration
    pub training_config: RLTrainingConfig,
}

/// Controller network for RL-based NAS
#[derive(Debug)]
pub struct ControllerNetwork<T: Float + Debug + Send + Sync + 'static> {
    /// Network weights
    pub weights: Vec<Array2<T>>,

    /// Network biases
    pub biases: Vec<Array1<T>>,

    /// Network architecture
    pub architecture: Vec<usize>,

    /// Activation functions
    pub activations: Vec<ActivationType>,

    /// Network type
    pub network_type: ControllerType,
}

/// Types of controller networks
#[derive(Debug, Clone, Copy)]
pub enum ControllerType {
    LSTM,
    Transformer,
    MLP,
    GRU,
}

/// Action space for architecture generation
#[derive(Debug, Clone)]
pub struct ActionSpace {
    /// Discrete actions
    pub discrete_actions: Vec<DiscreteAction>,

    /// Continuous actions
    pub continuous_actions: Vec<ContinuousAction>,

    /// Action constraints
    pub constraints: Vec<ActionConstraint>,

    /// Action embedding size
    pub embedding_size: usize,
}

/// Discrete action types
#[derive(Debug, Clone)]
pub struct DiscreteAction {
    /// Action name
    pub name: String,

    /// Possible values
    pub values: Vec<String>,

    /// Current probability distribution
    pub probabilities: Vec<f64>,
}

/// Continuous action types
#[derive(Debug, Clone)]
pub struct ContinuousAction {
    /// Action name
    pub name: String,

    /// Value range
    pub range: (f64, f64),

    /// Current mean and std
    pub mean: f64,
    pub std: f64,
}

/// Action constraints
#[derive(Debug, Clone)]
pub struct ActionConstraint {
    /// Constraint type
    pub constraint_type: ConstraintType,

    /// Constraint parameters
    pub parameters: HashMap<String, f64>,

    /// Constraint description
    pub description: String,
}

/// Types of action constraints
#[derive(Debug, Clone, Copy)]
pub enum ConstraintType {
    Dependency,    // Action depends on previous actions
    Exclusion,     // Action excludes other actions
    Range,         // Action value must be in range
    Conditional,   // Action only valid under conditions
}

/// State representation for the controller
#[derive(Debug)]
pub struct StateRepresentation<T: Float + Debug + Send + Sync + 'static> {
    /// Current architecture state
    pub architecture_state: Vec<T>,

    /// Performance history
    pub performance_history: Vec<T>,

    /// Resource usage state
    pub resource_state: Vec<T>,

    /// Search progress state
    pub progress_state: Vec<T>,

    /// State embedding
    pub embedding: Array1<T>,
}

/// Policy parameters
#[derive(Debug)]
pub struct PolicyParameters<T: Float + Debug + Send + Sync + 'static> {
    /// Policy type
    pub policy_type: PolicyType,

    /// Learning rate
    pub learning_rate: T,

    /// Discount factor
    pub gamma: T,

    /// Exploration parameters
    pub exploration: ExplorationParameters<T>,

    /// Entropy regularization weight
    pub entropy_weight: T,
}

/// Types of RL policies
#[derive(Debug, Clone, Copy)]
pub enum PolicyType {
    REINFORCE,
    ActorCritic,
    PPO,
    TRPO,
    SAC,
}

/// Exploration parameters
#[derive(Debug)]
pub struct ExplorationParameters<T: Float + Debug + Send + Sync + 'static> {
    /// Exploration strategy
    pub strategy: ExplorationStrategy,

    /// Initial exploration rate
    pub initial_rate: T,

    /// Final exploration rate
    pub final_rate: T,

    /// Decay schedule
    pub decay_schedule: DecaySchedule,

    /// Current exploration rate
    pub current_rate: T,
}

/// Exploration strategies
#[derive(Debug, Clone, Copy)]
pub enum ExplorationStrategy {
    EpsilonGreedy,
    BoltzmannExploration,
    UCB,
    ThompsonSampling,
    NoiseInjection,
}

/// Decay schedules for exploration
#[derive(Debug, Clone, Copy)]
pub enum DecaySchedule {
    Linear,
    Exponential,
    Polynomial,
    Cosine,
    Step,
}

/// Experience replay buffer
#[derive(Debug)]
pub struct ExperienceReplayBuffer<T: Float + Debug + Send + Sync + 'static> {
    /// Buffer capacity
    pub capacity: usize,

    /// Current size
    pub size: usize,

    /// Buffer index
    pub index: usize,

    /// States
    pub states: Vec<Vec<T>>,

    /// Actions
    pub actions: Vec<ActionSequence>,

    /// Rewards
    pub rewards: Vec<T>,

    /// Next states
    pub next_states: Vec<Vec<T>>,

    /// Done flags
    pub dones: Vec<bool>,
}

/// Action sequence for architecture generation
#[derive(Debug, Clone)]
pub struct ActionSequence {
    /// Discrete action indices
    pub discrete_actions: Vec<usize>,

    /// Continuous action values
    pub continuous_actions: Vec<f64>,

    /// Action probabilities
    pub probabilities: Vec<f64>,

    /// Action log probabilities
    pub log_probabilities: Vec<f64>,
}

/// RL training configuration
#[derive(Debug, Clone)]
pub struct RLTrainingConfig {
    /// Batch size for training
    pub batch_size: usize,

    /// Number of training epochs
    pub num_epochs: usize,

    /// Update frequency
    pub update_frequency: usize,

    /// Target network update frequency
    pub target_update_frequency: usize,

    /// Maximum episode length
    pub max_episode_length: usize,

    /// Reward normalization
    pub normalize_rewards: bool,

    /// Baseline subtraction
    pub use_baseline: bool,
}

impl Default for RLTrainingConfig {
    fn default() -> Self {
        Self {
            batch_size: 32,
            num_epochs: 10,
            update_frequency: 10,
            target_update_frequency: 100,
            max_episode_length: 50,
            normalize_rewards: true,
            use_baseline: true,
        }
    }
}

impl<T: Float + Debug + Default + std::fmt::Debug + Send + Sync> RLSearchState<T> {
    /// Create new RL search state
    pub fn new() -> Self {
        Self {
            controller: ControllerNetwork::new(ControllerType::LSTM),
            action_space: ActionSpace::default(),
            state_representation: StateRepresentation::new(),
            reward_history: VecDeque::with_capacity(1000),
            policy_parameters: PolicyParameters::default(),
            replay_buffer: ExperienceReplayBuffer::new(10000),
            training_config: RLTrainingConfig::default(),
        }
    }

    /// Generate architecture using the controller
    pub fn generate_architecture(
        &mut self,
        search_space: &super::super::space::ArchitectureSearchSpace,
    ) -> Result<ArchitectureCandidate, super::SearchError> {
        // Reset state
        self.state_representation.reset();

        // Generate action sequence
        let action_sequence = self.sample_action_sequence(search_space)?;

        // Convert action sequence to architecture
        let architecture = self.actions_to_architecture(&action_sequence, search_space)?;

        // Create candidate
        let candidate = ArchitectureCandidate::new(
            format!("rl_arch_{}", self.reward_history.len()),
            architecture,
        );

        Ok(candidate)
    }

    /// Sample action sequence from the controller
    fn sample_action_sequence(
        &mut self,
        _search_space: &super::super::space::ArchitectureSearchSpace,
    ) -> Result<ActionSequence, super::SearchError> {
        let mut discrete_actions = Vec::new();
        let mut continuous_actions = Vec::new();
        let mut probabilities = Vec::new();
        let mut log_probabilities = Vec::new();

        // Sample discrete actions
        for discrete_action in &self.action_space.discrete_actions {
            let action_idx = self.sample_discrete_action(discrete_action)?;
            discrete_actions.push(action_idx);
            
            if action_idx < discrete_action.probabilities.len() {
                let prob = discrete_action.probabilities[action_idx];
                probabilities.push(prob);
                log_probabilities.push(prob.ln());
            }
        }

        // Sample continuous actions
        for continuous_action in &self.action_space.continuous_actions {
            let action_value = self.sample_continuous_action(continuous_action)?;
            continuous_actions.push(action_value);
        }

        Ok(ActionSequence {
            discrete_actions,
            continuous_actions,
            probabilities,
            log_probabilities,
        })
    }

    /// Sample discrete action
    fn sample_discrete_action(&self, action: &DiscreteAction) -> Result<usize, super::SearchError> {
        let random_value: f64 = scirs2_core::random::random();
        let mut cumulative_prob = 0.0;

        for (idx, &prob) in action.probabilities.iter().enumerate() {
            cumulative_prob += prob;
            if random_value <= cumulative_prob {
                return Ok(idx);
            }
        }

        // Fallback to last action
        Ok(action.probabilities.len() - 1)
    }

    /// Sample continuous action
    fn sample_continuous_action(&self, action: &ContinuousAction) -> Result<f64, super::SearchError> {
        // Sample from normal distribution
        let normal_sample: f64 = scirs2_core::random::random::<f64>() * 2.0 - 1.0; // Simple approximation
        let value = action.mean + normal_sample * action.std;

        // Clamp to range
        Ok(value.max(action.range.0).min(action.range.1))
    }

    /// Convert action sequence to architecture
    fn actions_to_architecture(
        &self,
        actions: &ActionSequence,
        search_space: &super::super::space::ArchitectureSearchSpace,
    ) -> Result<ArchitectureSpec, super::SearchError> {
        use super::super::architecture::{LayerSpec, LayerDimensions, GlobalArchitectureConfig};

        // Simplified conversion
        let num_layers = if !actions.discrete_actions.is_empty() {
            (actions.discrete_actions[0] % 5) + 1  // 1-5 layers
        } else {
            3
        };

        let mut layers = Vec::new();
        for i in 0..num_layers {
            let layer_type_idx = if i < actions.discrete_actions.len() {
                actions.discrete_actions[i] % search_space.layer_types.len()
            } else {
                0
            };

            let layer_type = search_space.layer_types[layer_type_idx];

            let hidden_size_idx = if i < actions.continuous_actions.len() {
                (actions.continuous_actions[i] * search_space.hidden_sizes.len() as f64) as usize
            } else {
                0
            };
            let hidden_size = search_space.hidden_sizes[hidden_size_idx.min(search_space.hidden_sizes.len() - 1)];

            let activation_idx = if i + 1 < actions.discrete_actions.len() {
                actions.discrete_actions[i + 1] % search_space.activation_functions.len()
            } else {
                0
            };
            let activation = search_space.activation_functions[activation_idx];

            let dimensions = LayerDimensions {
                input_dim: if i == 0 { hidden_size } else { layers[i-1].dimensions.output_dim },
                output_dim: hidden_size,
                hidden_dims: vec![],
            };

            layers.push(LayerSpec::new(layer_type, dimensions, activation));
        }

        Ok(ArchitectureSpec::new(layers, GlobalArchitectureConfig::default()))
    }

    /// Update controller with reward
    pub fn update_with_reward(&mut self, reward: f64) {
        self.reward_history.push_back(reward);
        
        // Keep history bounded
        if self.reward_history.len() > 1000 {
            self.reward_history.pop_front();
        }

        // Update policy (simplified)
        if self.reward_history.len() % self.training_config.update_frequency == 0 {
            let _ = self.update_policy();
        }
    }

    /// Update policy using collected experiences
    fn update_policy(&mut self) -> Result<(), super::SearchError> {
        match self.policy_parameters.policy_type {
            PolicyType::REINFORCE => {
                self.update_reinforce()
            }
            PolicyType::ActorCritic => {
                self.update_actor_critic()
            }
            PolicyType::PPO => {
                self.update_ppo()
            }
            _ => {
                // Default to simple gradient ascent
                Ok(())
            }
        }
    }

    /// Update using REINFORCE
    fn update_reinforce(&mut self) -> Result<(), super::SearchError> {
        // Simplified REINFORCE update
        // In practice, would compute policy gradients and update controller weights
        Ok(())
    }

    /// Update using Actor-Critic
    fn update_actor_critic(&mut self) -> Result<(), super::SearchError> {
        // Simplified Actor-Critic update
        // Would update both actor and critic networks
        Ok(())
    }

    /// Update using PPO
    fn update_ppo(&mut self) -> Result<(), super::SearchError> {
        // Simplified PPO update
        // Would implement clipped surrogate objective
        Ok(())
    }

    /// Check if training should stop
    pub fn should_stop_training(&self) -> bool {
        // Stop if converged (no improvement in recent rewards)
        if self.reward_history.len() < 100 {
            return false;
        }

        let recent_mean = self.reward_history
            .iter()
            .rev()
            .take(50)
            .sum::<f64>() / 50.0;

        let older_mean = self.reward_history
            .iter()
            .rev()
            .skip(50)
            .take(50)
            .sum::<f64>() / 50.0;

        (recent_mean - older_mean).abs() < 0.001
    }
}

impl<T: Float + Debug + Default + Send + Sync> ControllerNetwork<T> {
    /// Create new controller network
    pub fn new(network_type: ControllerType) -> Self {
        let architecture = match network_type {
            ControllerType::LSTM => vec![128, 64, 32],
            ControllerType::Transformer => vec![256, 128, 64],
            ControllerType::MLP => vec![128, 64, 32, 16],
            ControllerType::GRU => vec![128, 64, 32],
        };

        let weights = Vec::new(); // Would initialize with random weights
        let biases = Vec::new();  // Would initialize with zeros
        let activations = vec![ActivationType::ReLU; architecture.len()];

        Self {
            weights,
            biases,
            architecture,
            activations,
            network_type,
        }
    }

    /// Forward pass through the network
    pub fn forward(&self, input: &[T]) -> Result<Vec<T>, super::SearchError> {
        // Simplified forward pass
        // In practice would implement actual network computation
        Ok(input.to_vec())
    }
}

impl Default for ActionSpace {
    fn default() -> Self {
        let discrete_actions = vec![
            DiscreteAction {
                name: "layer_type".to_string(),
                values: vec!["linear".to_string(), "lstm".to_string(), "attention".to_string()],
                probabilities: vec![0.33, 0.33, 0.34],
            },
            DiscreteAction {
                name: "activation".to_string(),
                values: vec!["relu".to_string(), "gelu".to_string(), "tanh".to_string()],
                probabilities: vec![0.33, 0.33, 0.34],
            },
        ];

        let continuous_actions = vec![
            ContinuousAction {
                name: "hidden_size".to_string(),
                range: (32.0, 512.0),
                mean: 128.0,
                std: 64.0,
            },
            ContinuousAction {
                name: "dropout_rate".to_string(),
                range: (0.0, 0.5),
                mean: 0.1,
                std: 0.05,
            },
        ];

        Self {
            discrete_actions,
            continuous_actions,
            constraints: Vec::new(),
            embedding_size: 64,
        }
    }
}

impl<T: Float + Debug + Default + Send + Sync> StateRepresentation<T> {
    /// Create new state representation
    pub fn new() -> Self {
        Self {
            architecture_state: Vec::new(),
            performance_history: Vec::new(),
            resource_state: Vec::new(),
            progress_state: Vec::new(),
            embedding: Array1::zeros(64),
        }
    }

    /// Reset state representation
    pub fn reset(&mut self) {
        self.architecture_state.clear();
        self.performance_history.clear();
        self.resource_state.clear();
        self.progress_state.clear();
        self.embedding.fill(T::zero());
    }

    /// Update state with new information
    pub fn update(&mut self, architecture: &ArchitectureSpec, performance: f64) {
        // Update architecture state
        self.architecture_state = vec![
            T::from(architecture.layers.len()).expect("unwrap failed"),
            T::from(architecture.parameter_count()).expect("unwrap failed"),
        ];

        // Update performance history
        self.performance_history.push(scirs2_core::numeric::NumCast::from(performance).unwrap_or_else(|| T::zero()));
        if self.performance_history.len() > 10 {
            self.performance_history.remove(0);
        }

        // Update embedding (simplified)
        for (i, value) in self.architecture_state.iter().enumerate() {
            if i < self.embedding.len() {
                self.embedding[i] = *value;
            }
        }
    }
}

impl<T: Float + Debug + Default + Send + Sync> PolicyParameters<T> {
    /// Create default policy parameters
    pub fn default() -> Self {
        Self {
            policy_type: PolicyType::REINFORCE,
            learning_rate: scirs2_core::numeric::NumCast::from(0.001).unwrap_or_else(|| T::zero()),
            gamma: scirs2_core::numeric::NumCast::from(0.99).unwrap_or_else(|| T::zero()),
            exploration: ExplorationParameters::default(),
            entropy_weight: scirs2_core::numeric::NumCast::from(0.01).unwrap_or_else(|| T::zero()),
        }
    }
}

impl<T: Float + Debug + Default + Send + Sync> ExplorationParameters<T> {
    /// Create default exploration parameters
    pub fn default() -> Self {
        Self {
            strategy: ExplorationStrategy::EpsilonGreedy,
            initial_rate: scirs2_core::numeric::NumCast::from(1.0).unwrap_or_else(|| T::zero()),
            final_rate: scirs2_core::numeric::NumCast::from(0.01).unwrap_or_else(|| T::zero()),
            decay_schedule: DecaySchedule::Exponential,
            current_rate: scirs2_core::numeric::NumCast::from(1.0).unwrap_or_else(|| T::zero()),
        }
    }

    /// Update exploration rate
    pub fn update_rate(&mut self, step: usize, total_steps: usize) {
        let progress = scirs2_core::numeric::NumCast::from(step as f64 / total_steps as f64).unwrap_or_else(|| T::zero());
        
        self.current_rate = match self.decay_schedule {
            DecaySchedule::Linear => {
                self.initial_rate - (self.initial_rate - self.final_rate) * progress
            }
            DecaySchedule::Exponential => {
                self.initial_rate * (self.final_rate / self.initial_rate).powf(progress)
            }
            _ => self.initial_rate,
        };
    }
}

impl<T: Float + Debug + Default + Send + Sync> ExperienceReplayBuffer<T> {
    /// Create new experience replay buffer
    pub fn new(capacity: usize) -> Self {
        Self {
            capacity,
            size: 0,
            index: 0,
            states: Vec::with_capacity(capacity),
            actions: Vec::with_capacity(capacity),
            rewards: Vec::with_capacity(capacity),
            next_states: Vec::with_capacity(capacity),
            dones: Vec::with_capacity(capacity),
        }
    }

    /// Add experience to buffer
    pub fn add_experience(
        &mut self,
        state: Vec<T>,
        action: ActionSequence,
        reward: T,
        next_state: Vec<T>,
        done: bool,
    ) {
        if self.size < self.capacity {
            self.states.push(state);
            self.actions.push(action);
            self.rewards.push(reward);
            self.next_states.push(next_state);
            self.dones.push(done);
            self.size += 1;
        } else {
            self.states[self.index] = state;
            self.actions[self.index] = action;
            self.rewards[self.index] = reward;
            self.next_states[self.index] = next_state;
            self.dones[self.index] = done;
        }

        self.index = (self.index + 1) % self.capacity;
    }

    /// Sample batch from buffer
    pub fn sample_batch(&self, batch_size: usize) -> Vec<usize> {
        let mut indices = Vec::new();
        
        for _ in 0..batch_size {
            let idx = scirs2_core::random::random::<usize>() % self.size;
            indices.push(idx);
        }
        
        indices
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rl_state_creation() {
        let state: RLSearchState<f64> = RLSearchState::new();
        assert_eq!(state.reward_history.len(), 0);
        assert!(matches!(state.controller.network_type, ControllerType::LSTM));
    }

    #[test]
    fn test_controller_network() {
        let controller: ControllerNetwork<f64> = ControllerNetwork::new(ControllerType::MLP);
        assert!(matches!(controller.network_type, ControllerType::MLP));
        assert!(!controller.architecture.is_empty());
    }

    #[test]
    fn test_action_space_default() {
        let action_space = ActionSpace::default();
        assert!(!action_space.discrete_actions.is_empty());
        assert!(!action_space.continuous_actions.is_empty());
    }

    #[test]
    fn test_experience_replay_buffer() {
        let mut buffer: ExperienceReplayBuffer<f64> = ExperienceReplayBuffer::new(100);
        
        buffer.add_experience(
            vec![1.0, 2.0],
            ActionSequence {
                discrete_actions: vec![0, 1],
                continuous_actions: vec![0.5],
                probabilities: vec![0.5, 0.3],
                log_probabilities: vec![-0.693, -1.204],
            },
            1.0,
            vec![2.0, 3.0],
            false,
        );
        
        assert_eq!(buffer.size, 1);
    }

    #[test]
    fn test_exploration_parameters() {
        let mut exploration: ExplorationParameters<f64> = ExplorationParameters::default();
        exploration.update_rate(50, 100);
        assert!(exploration.current_rate < exploration.initial_rate);
    }
}