optirs-learned 0.3.1

use std::fmt::Debug;
// Neural architecture search strategies and algorithms
//
// This module implements various search strategies for neural architecture search,
// including evolutionary algorithms, Bayesian optimization, reinforcement learning,
// and other advanced search methods.

pub mod strategies;
pub mod evolutionary;
pub mod bayesian;
pub mod reinforcement;
pub mod differentiable;
pub mod progressive;

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

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

pub use strategies::*;
pub use evolutionary::*;
pub use bayesian::*;
pub use reinforcement::*;
pub use differentiable::*;
pub use progressive::*;

/// Search strategy implementation
pub struct SearchStrategy<T: Float + Debug + Send + Sync + 'static> {
    /// Strategy type
    strategy_type: SearchStrategyType,

    /// Random number generator
    #[allow(dead_code)]
    rng: Box<dyn scirs2_core::random::RngCore + Send>,

    /// Strategy-specific state
    state: SearchStrategyState<T>,

    /// Optimization history
    optimization_history: Vec<OptimizationStep<T>>,

    /// Current best architectures
    best_architectures: Vec<ArchitectureCandidate>,
}

impl<T: Float + Debug + std::fmt::Debug + Send + Sync> std::fmt::Debug for SearchStrategy<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SearchStrategy")
            .field("strategy_type", &self.strategy_type)
            .field("state", &self.state)
            .field("optimization_history", &self.optimization_history)
            .field("best_architectures", &self.best_architectures)
            .finish()
    }
}

/// Types of search strategies
#[derive(Debug, Clone, Copy, Default)]
pub enum SearchStrategyType {
    /// Random search
    #[default]
    Random,

    /// Evolutionary algorithm
    Evolutionary,

    /// Bayesian optimization
    BayesianOptimization,

    /// Reinforcement learning
    ReinforcementLearning,

    /// Differentiable NAS
    DifferentiableNAS,

    /// Progressive search
    Progressive,

    /// Multi-objective search
    MultiObjective,

    /// Hyperband-based search
    Hyperband,
}

/// Search strategy state
#[derive(Debug)]
pub enum SearchStrategyState<T: Float + Debug + Send + Sync + 'static> {
    Random(RandomSearchState),
    Evolutionary(EvolutionarySearchState<T>),
    Bayesian(BayesianOptimizationState<T>),
    ReinforcementLearning(RLSearchState<T>),
    Differentiable(DifferentiableNASState<T>),
    Progressive(ProgressiveSearchState<T>),
    MultiObjective(MultiObjectiveState<T>),
}

/// Random search state
#[derive(Debug, Default)]
pub struct RandomSearchState {
    /// Sampling budget remaining
    pub budget_remaining: usize,

    /// Sampling history
    pub sampling_history: Vec<String>,
}

/// Optimization step record
#[derive(Debug, Clone)]
pub struct OptimizationStep<T: Float + Debug + Send + Sync + 'static> {
    /// Step number
    pub step: usize,

    /// Architecture evaluated
    pub architecture: ArchitectureSpec,

    /// Performance achieved
    pub performance: T,

    /// Step timestamp
    pub timestamp: std::time::Instant,

    /// Step type
    pub step_type: OptimizationStepType,

    /// Additional metadata
    pub metadata: HashMap<String, String>,
}

/// Types of optimization steps
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum OptimizationStepType {
    Random,
    Mutation,
    Crossover,
    Selection,
    Evaluation,
    Acquisition,
    PolicyUpdate,
    ArchitectureUpdate,
}

// Implementation methods
impl<T: Float + Debug + Default + std::fmt::Debug + Send + Sync> SearchStrategy<T> {
    /// Create new search strategy
    pub fn new(strategy_type: SearchStrategyType, config: SearchConfig) -> Self {
        let rng = Box::new(scirs2_core::random::thread_rng());
        
        let state = match strategy_type {
            SearchStrategyType::Random => {
                SearchStrategyState::Random(RandomSearchState {
                    budget_remaining: config.budget,
                    sampling_history: Vec::new(),
                })
            }
            SearchStrategyType::Evolutionary => {
                SearchStrategyState::Evolutionary(EvolutionarySearchState::new(config.population_size))
            }
            SearchStrategyType::BayesianOptimization => {
                SearchStrategyState::Bayesian(BayesianOptimizationState::new())
            }
            SearchStrategyType::ReinforcementLearning => {
                SearchStrategyState::ReinforcementLearning(RLSearchState::new())
            }
            SearchStrategyType::DifferentiableNAS => {
                SearchStrategyState::Differentiable(DifferentiableNASState::new())
            }
            SearchStrategyType::Progressive => {
                SearchStrategyState::Progressive(ProgressiveSearchState::new())
            }
            SearchStrategyType::MultiObjective => {
                SearchStrategyState::MultiObjective(MultiObjectiveState::new())
            }
            SearchStrategyType::Hyperband => {
                // Default to progressive for now
                SearchStrategyState::Progressive(ProgressiveSearchState::new())
            }
        };

        Self {
            strategy_type,
            rng,
            state,
            optimization_history: Vec::new(),
            best_architectures: Vec::new(),
        }
    }

    /// Generate next architecture candidate
    pub fn generate_candidate(&mut self) -> Result<ArchitectureCandidate, SearchError> {
        match &mut self.state {
            SearchStrategyState::Random(state) => {
                self.generate_random_candidate(state)
            }
            SearchStrategyState::Evolutionary(state) => {
                self.generate_evolutionary_candidate(state)
            }
            SearchStrategyState::Bayesian(state) => {
                self.generate_bayesian_candidate(state)
            }
            SearchStrategyState::ReinforcementLearning(state) => {
                self.generate_rl_candidate(state)
            }
            SearchStrategyState::Differentiable(state) => {
                self.generate_differentiable_candidate(state)
            }
            SearchStrategyState::Progressive(state) => {
                self.generate_progressive_candidate(state)
            }
            SearchStrategyState::MultiObjective(state) => {
                self.generate_multiobjective_candidate(state)
            }
        }
    }

    /// Update strategy with evaluation result
    pub fn update_with_result(&mut self, candidate: &ArchitectureCandidate, performance: T) {
        // Record optimization step
        let step = OptimizationStep {
            step: self.optimization_history.len(),
            architecture: candidate.architecture.clone(),
            performance,
            timestamp: std::time::Instant::now(),
            step_type: OptimizationStepType::Evaluation,
            metadata: HashMap::new(),
        };
        
        self.optimization_history.push(step);

        // Update best architectures
        self.update_best_architectures(candidate.clone());

        // Strategy-specific updates
        match &mut self.state {
            SearchStrategyState::Evolutionary(state) => {
                self.update_evolutionary_state(state, candidate, performance);
            }
            SearchStrategyState::Bayesian(state) => {
                self.update_bayesian_state(state, candidate, performance);
            }
            SearchStrategyState::ReinforcementLearning(state) => {
                self.update_rl_state(state, candidate, performance);
            }
            SearchStrategyState::Differentiable(state) => {
                self.update_differentiable_state(state, candidate, performance);
            }
            SearchStrategyState::Progressive(state) => {
                self.update_progressive_state(state, candidate, performance);
            }
            SearchStrategyState::MultiObjective(state) => {
                self.update_multiobjective_state(state, candidate, performance);
            }
            _ => {} // No updates needed for random search
        }
    }

    /// Get current best architectures
    pub fn get_best_architectures(&self) -> &Vec<ArchitectureCandidate> {
        &self.best_architectures
    }

    /// Get optimization history
    pub fn get_optimization_history(&self) -> &Vec<OptimizationStep<T>> {
        &self.optimization_history
    }

    /// Check if search should terminate
    pub fn should_terminate(&self) -> bool {
        match &self.state {
            SearchStrategyState::Random(state) => {
                state.budget_remaining == 0
            }
            SearchStrategyState::Evolutionary(state) => {
                state.generation >= 100 // Default max generations
            }
            _ => false // Continue indefinitely for other strategies
        }
    }

    // Strategy-specific generation methods
    fn generate_random_candidate(&mut self, state: &mut RandomSearchState) -> Result<ArchitectureCandidate, SearchError> {
        if state.budget_remaining == 0 {
            return Err(SearchError::BudgetExhausted);
        }

        state.budget_remaining -= 1;
        
        // Generate random architecture (simplified)
        let id = format!("random_{}", state.sampling_history.len());
        let arch_spec = ArchitectureSpec::new(vec![], super::architecture::GlobalArchitectureConfig::default());
        let candidate = ArchitectureCandidate::new(id.clone(), arch_spec);
        
        state.sampling_history.push(id);
        
        Ok(candidate)
    }

    fn generate_evolutionary_candidate(&mut self, _state: &mut EvolutionarySearchState<T>) -> Result<ArchitectureCandidate, SearchError> {
        // Placeholder - detailed implementation in evolutionary module
        Err(SearchError::NotImplemented("Evolutionary generation".to_string()))
    }

    fn generate_bayesian_candidate(&mut self, _state: &mut BayesianOptimizationState<T>) -> Result<ArchitectureCandidate, SearchError> {
        // Placeholder - detailed implementation in bayesian module
        Err(SearchError::NotImplemented("Bayesian generation".to_string()))
    }

    fn generate_rl_candidate(&mut self, _state: &mut RLSearchState<T>) -> Result<ArchitectureCandidate, SearchError> {
        // Placeholder - detailed implementation in reinforcement module
        Err(SearchError::NotImplemented("RL generation".to_string()))
    }

    fn generate_differentiable_candidate(&mut self, _state: &mut DifferentiableNASState<T>) -> Result<ArchitectureCandidate, SearchError> {
        // Placeholder - detailed implementation in differentiable module
        Err(SearchError::NotImplemented("Differentiable generation".to_string()))
    }

    fn generate_progressive_candidate(&mut self, _state: &mut ProgressiveSearchState<T>) -> Result<ArchitectureCandidate, SearchError> {
        // Placeholder - detailed implementation in progressive module
        Err(SearchError::NotImplemented("Progressive generation".to_string()))
    }

    fn generate_multiobjective_candidate(&mut self, _state: &mut MultiObjectiveState<T>) -> Result<ArchitectureCandidate, SearchError> {
        // Placeholder - detailed implementation in multi-objective module
        Err(SearchError::NotImplemented("Multi-objective generation".to_string()))
    }

    // Strategy-specific update methods
    fn update_evolutionary_state(&mut self, _state: &mut EvolutionarySearchState<T>, _candidate: &ArchitectureCandidate, _performance: T) {
        // Implementation in evolutionary module
    }

    fn update_bayesian_state(&mut self, _state: &mut BayesianOptimizationState<T>, _candidate: &ArchitectureCandidate, _performance: T) {
        // Implementation in bayesian module
    }

    fn update_rl_state(&mut self, _state: &mut RLSearchState<T>, _candidate: &ArchitectureCandidate, _performance: T) {
        // Implementation in reinforcement module
    }

    fn update_differentiable_state(&mut self, _state: &mut DifferentiableNASState<T>, _candidate: &ArchitectureCandidate, _performance: T) {
        // Implementation in differentiable module
    }

    fn update_progressive_state(&mut self, _state: &mut ProgressiveSearchState<T>, _candidate: &ArchitectureCandidate, _performance: T) {
        // Implementation in progressive module
    }

    fn update_multiobjective_state(&mut self, _state: &mut MultiObjectiveState<T>, _candidate: &ArchitectureCandidate, _performance: T) {
        // Implementation in multi-objective module
    }

    fn update_best_architectures(&mut self, candidate: ArchitectureCandidate) {
        self.best_architectures.push(candidate);
        
        // Keep only top N architectures
        self.best_architectures.sort_by(|a, b| {
            b.performance.optimization_performance
                .partial_cmp(&a.performance.optimization_performance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        
        if self.best_architectures.len() > 10 {
            self.best_architectures.truncate(10);
        }
    }
}

/// Search configuration
#[derive(Debug, Clone)]
pub struct SearchConfig {
    /// Search budget (number of evaluations)
    pub budget: usize,

    /// Population size (for evolutionary strategies)
    pub population_size: usize,

    /// Maximum iterations
    pub max_iterations: usize,

    /// Convergence threshold
    pub convergence_threshold: f64,

    /// Early stopping patience
    pub early_stopping_patience: usize,

    /// Random seed
    pub random_seed: Option<u64>,

    /// Strategy-specific parameters
    pub strategy_params: HashMap<String, f64>,
}

impl Default for SearchConfig {
    fn default() -> Self {
        Self {
            budget: 1000,
            population_size: 50,
            max_iterations: 100,
            convergence_threshold: 0.001,
            early_stopping_patience: 10,
            random_seed: None,
            strategy_params: HashMap::new(),
        }
    }
}

/// Search errors
#[derive(Debug, Clone)]
pub enum SearchError {
    BudgetExhausted,
    InvalidConfiguration(String),
    GenerationFailed(String),
    EvaluationFailed(String),
    NotImplemented(String),
    StrategyError(String),
}

impl std::fmt::Display for SearchError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SearchError::BudgetExhausted => write!(f, "Search budget exhausted"),
            SearchError::InvalidConfiguration(msg) => write!(f, "Invalid configuration: {}", msg),
            SearchError::GenerationFailed(msg) => write!(f, "Generation failed: {}", msg),
            SearchError::EvaluationFailed(msg) => write!(f, "Evaluation failed: {}", msg),
            SearchError::NotImplemented(msg) => write!(f, "Not implemented: {}", msg),
            SearchError::StrategyError(msg) => write!(f, "Strategy error: {}", msg),
        }
    }
}

impl std::error::Error for SearchError {}

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

    #[test]
    fn test_search_strategy_creation() {
        let config = SearchConfig::default();
        let strategy: SearchStrategy<f64> = SearchStrategy::new(SearchStrategyType::Random, config);
        
        match strategy.state {
            SearchStrategyState::Random(_) => {}
            _ => panic!("Expected random search state"),
        }
    }

    #[test]
    fn test_random_candidate_generation() {
        let config = SearchConfig { budget: 5, ..Default::default() };
        let mut strategy: SearchStrategy<f64> = SearchStrategy::new(SearchStrategyType::Random, config);
        
        for _ in 0..5 {
            let result = strategy.generate_candidate();
            assert!(result.is_ok());
        }
        
        // Should fail when budget exhausted
        let result = strategy.generate_candidate();
        assert!(matches!(result, Err(SearchError::BudgetExhausted)));
    }

    #[test]
    fn test_termination_condition() {
        let config = SearchConfig { budget: 0, ..Default::default() };
        let strategy: SearchStrategy<f64> = SearchStrategy::new(SearchStrategyType::Random, config);
        
        assert!(strategy.should_terminate());
    }

    #[test]
    fn test_best_architecture_tracking() {
        let config = SearchConfig::default();
        let mut strategy: SearchStrategy<f64> = SearchStrategy::new(SearchStrategyType::Random, config);
        
        let arch = ArchitectureSpec::new(vec![], super::architecture::GlobalArchitectureConfig::default());
        let candidate = ArchitectureCandidate::new("test".to_string(), arch);
        
        strategy.update_with_result(&candidate, 0.8);
        assert_eq!(strategy.get_best_architectures().len(), 1);
    }
}