optirs-nas 0.3.1

// Neural predictor-based search strategy

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

#[allow(unused_imports)]
use crate::error::Result;
use crate::nas_engine::{OptimizerArchitecture, SearchResult, SearchSpaceConfig};
use crate::EvaluationMetric;

use super::random::RandomSearch;
use super::{SearchStrategy, SearchStrategyStatistics};

/// Neural predictor-based search
pub struct NeuralPredictorSearch<T: Float + Debug + Send + Sync + 'static> {
    predictor_network: PredictorNetwork<T>,
    architecture_encoder: ArchitectureEncoder<T>,
    search_optimizer: SearchOptimizer<T>,
    confidence_threshold: T,
    statistics: SearchStrategyStatistics<T>,
    uncertainty_sampling: bool,
}

/// Predictor network for neural predictor search
#[derive(Debug)]
pub struct PredictorNetwork<T: Float + Debug + Send + Sync + 'static> {
    layers: Vec<PredictorLayer<T>>,
    dropout_rates: Vec<T>,
    architecture: Vec<usize>,
}

/// Predictor layer
#[derive(Debug)]
pub struct PredictorLayer<T: Float + Debug + Send + Sync + 'static> {
    weights: Array2<T>,
    bias: Array1<T>,
    activation: ActivationFunction,
}

/// Activation functions
#[derive(Debug, Clone, Copy)]
pub enum ActivationFunction {
    ReLU,
    GELU,
    Swish,
    Tanh,
    Sigmoid,
}

/// Architecture encoder for neural predictor
#[derive(Debug)]
pub struct ArchitectureEncoder<T: Float + Debug + Send + Sync + 'static> {
    encoding_weights: Array2<T>,
    _embeddingdim: usize,
    max_components: usize,
}

/// Search optimizer for neural predictor
#[derive(Debug)]
pub struct SearchOptimizer<T: Float + Debug + Send + Sync + 'static> {
    optimizer_type: SearchOptimizerType,
    _learningrate: T,
    momentum: T,
    parameters: HashMap<String, Array1<T>>,
}

/// Search optimizer types
#[derive(Debug, Clone, Copy)]
pub enum SearchOptimizerType {
    Adam,
    SGD,
    RMSprop,
    AdamW,
}

impl<
        T: Float + Debug + Default + Clone + Send + Sync + std::fmt::Debug + 'static + std::iter::Sum,
    > NeuralPredictorSearch<T>
{
    pub fn new(
        predictor_architecture: Vec<usize>,
        embeddingdim: usize,
        confidence_threshold: f64,
    ) -> Self {
        Self {
            predictor_network: PredictorNetwork::new(predictor_architecture),
            architecture_encoder: ArchitectureEncoder::new(embeddingdim),
            search_optimizer: SearchOptimizer::new(
                SearchOptimizerType::Adam,
                scirs2_core::numeric::NumCast::from(0.001).unwrap_or_else(|| T::zero()),
            ),
            confidence_threshold: scirs2_core::numeric::NumCast::from(confidence_threshold)
                .unwrap_or_else(|| T::zero()),
            statistics: SearchStrategyStatistics::default(),
            uncertainty_sampling: true,
        }
    }

    fn predict_performance(&self, architecture: &OptimizerArchitecture<T>) -> Result<(T, T)> {
        // Encode architecture
        let encoded = self.architecture_encoder.encode(architecture)?;

        // Forward pass through predictor network
        let (prediction, uncertainty) =
            self.predictor_network.forward_with_uncertainty(&encoded)?;

        Ok((prediction, uncertainty))
    }

    fn train_predictor(
        &mut self,
        architectures: &[OptimizerArchitecture<T>],
        performances: &[T],
    ) -> Result<()> {
        if architectures.len() != performances.len() || architectures.is_empty() {
            return Ok(());
        }

        // Encode all architectures
        let encoded_archs: std::result::Result<Vec<_>, _> = architectures
            .iter()
            .map(|arch| self.architecture_encoder.encode(arch))
            .collect();
        let encoded_archs = encoded_archs?;

        // Train predictor network
        for (encoded_arch, &target_performance) in encoded_archs.iter().zip(performances.iter()) {
            let (_prediction, _) = self
                .predictor_network
                .forward_with_uncertainty(encoded_arch)?;

            // Simplified gradient update
            self.predictor_network.backward_update(
                encoded_arch,
                target_performance,
                &mut self.search_optimizer,
            )?;
        }

        Ok(())
    }

    fn generate_candidate_with_uncertainty(
        &mut self,
        searchspace: &SearchSpaceConfig,
    ) -> Result<OptimizerArchitecture<T>> {
        // Generate multiple candidates and select based on uncertainty
        let num_candidates = 50;
        let mut candidates = Vec::new();
        let mut random_search = RandomSearch::<T>::new(None);
        random_search.initialize(searchspace)?;

        for _ in 0..num_candidates {
            candidates.push(random_search.generate_architecture(searchspace, &VecDeque::new())?);
        }

        // Select candidate with highest uncertainty (for exploration) or highest predicted
        // performance (for exploitation)
        let mut best_candidate = candidates[0].clone();
        let mut best_score = T::neg_infinity();

        for candidate in candidates {
            let (predicted_perf, uncertainty) = self.predict_performance(&candidate)?;

            // Combine prediction and uncertainty for selection
            let score = if self.uncertainty_sampling {
                predicted_perf + uncertainty // UCB-style selection
            } else {
                predicted_perf // Pure exploitation
            };

            if score > best_score {
                best_score = score;
                best_candidate = candidate;
            }
        }

        Ok(best_candidate)
    }
}

impl<
        T: Float + Debug + Default + Clone + Send + Sync + std::fmt::Debug + 'static + std::iter::Sum,
    > SearchStrategy<T> for NeuralPredictorSearch<T>
{
    fn initialize(&mut self, _searchspace: &SearchSpaceConfig) -> Result<()> {
        // Initialize predictor network with random weights
        self.predictor_network.initialize()?;
        Ok(())
    }

    fn generate_architecture(
        &mut self,
        searchspace: &SearchSpaceConfig,
        history: &VecDeque<SearchResult<T>>,
    ) -> Result<OptimizerArchitecture<T>> {
        // Train predictor if enough data is available
        if history.len() > 10 {
            let architectures: Vec<_> = history.iter().map(|r| r.architecture.clone()).collect();
            let performances: Vec<_> = history
                .iter()
                .filter_map(|r| {
                    r.evaluation_results
                        .metric_scores
                        .get(&EvaluationMetric::FinalPerformance)
                })
                .cloned()
                .collect();

            if architectures.len() == performances.len() {
                self.train_predictor(&architectures, &performances)?;
            }
        }

        // Generate candidate based on predictor
        let architecture = if history.len() > 5 {
            self.generate_candidate_with_uncertainty(searchspace)?
        } else {
            // Use random search for initial exploration
            let mut random_search = RandomSearch::<T>::new(None);
            random_search.initialize(searchspace)?;
            random_search.generate_architecture(searchspace, history)?
        };

        self.statistics.total_architectures_generated += 1;
        Ok(architecture)
    }

    fn update_with_results(&mut self, results: &[SearchResult<T>]) -> Result<()> {
        if results.is_empty() {
            return Ok(());
        }

        // Extract architectures and performances
        let architectures: Vec<_> = results.iter().map(|r| r.architecture.clone()).collect();
        let performances: Vec<_> = results
            .iter()
            .filter_map(|r| {
                r.evaluation_results
                    .metric_scores
                    .get(&EvaluationMetric::FinalPerformance)
            })
            .cloned()
            .collect();

        if architectures.len() == performances.len() && !performances.is_empty() {
            // Update statistics
            self.statistics.best_performance = performances
                .iter()
                .max_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
                .cloned()
                .unwrap_or(T::zero());

            let sum: T = performances.iter().cloned().sum();
            self.statistics.average_performance =
                sum / T::from(performances.len()).expect("conversion from usize to T failed");

            // Train predictor with new data
            self.train_predictor(&architectures, &performances)?;
        }

        Ok(())
    }

    fn name(&self) -> &str {
        "NeuralPredictorSearch"
    }

    fn get_statistics(&self) -> SearchStrategyStatistics<T> {
        let mut stats = self.statistics.clone();
        stats.exploration_rate = if self.uncertainty_sampling {
            scirs2_core::numeric::NumCast::from(0.7).unwrap_or_else(|| T::zero())
        } else {
            scirs2_core::numeric::NumCast::from(0.3).unwrap_or_else(|| T::zero())
        };
        stats.exploitation_rate = T::one() - stats.exploration_rate;
        stats
    }
}

// Implementation for supporting components
impl<T: Float + Debug + Default + Clone + 'static + std::iter::Sum + Send + Sync>
    PredictorNetwork<T>
{
    fn new(architecture: Vec<usize>) -> Self {
        let mut layers = Vec::new();
        for i in 0..architecture.len() - 1 {
            layers.push(PredictorLayer::new(architecture[i], architecture[i + 1]));
        }

        Self {
            layers,
            dropout_rates: vec![
                scirs2_core::numeric::NumCast::from(0.1)
                    .unwrap_or_else(|| T::zero());
                architecture.len() - 1
            ],
            architecture,
        }
    }

    fn initialize(&mut self) -> Result<()> {
        for layer in &mut self.layers {
            layer.initialize()?;
        }
        Ok(())
    }

    fn forward_with_uncertainty(&self, input: &Array1<T>) -> Result<(T, T)> {
        let mut current = input.clone();

        // Forward pass through all layers
        for (i, layer) in self.layers.iter().enumerate() {
            current = layer.forward(&current)?;

            // Apply dropout for uncertainty estimation (Monte Carlo dropout)
            if i < self.dropout_rates.len() {
                current = self.apply_dropout(&current, self.dropout_rates[i]);
            }
        }

        // For simplicity, return the first output as prediction and a simple uncertainty estimate
        let prediction = current[0];
        let uncertainty = current.iter().map(|&x| x * x).sum::<T>().sqrt()
            * scirs2_core::numeric::NumCast::from(0.1).unwrap_or_else(|| T::zero());

        Ok((prediction, uncertainty))
    }

    fn backward_update(
        &mut self,
        _input: &Array1<T>,
        _target: T,
        _optimizer: &mut SearchOptimizer<T>,
    ) -> Result<()> {
        // Simplified backward pass - in practice would implement proper backpropagation
        Ok(())
    }

    fn apply_dropout(&self, input: &Array1<T>, dropoutrate: T) -> Array1<T> {
        input.mapv(|x| {
            if scirs2_core::random::Random::default().random::<f64>()
                < dropoutrate.to_f64().unwrap_or(0.0)
            {
                T::zero()
            } else {
                x / (T::one() - dropoutrate)
            }
        })
    }
}

impl<T: Float + Debug + Default + Clone + 'static + Send + Sync> PredictorLayer<T> {
    fn new(input_size: usize, outputsize: usize) -> Self {
        Self {
            weights: Array2::zeros((outputsize, input_size)),
            bias: Array1::zeros(outputsize),
            activation: ActivationFunction::ReLU,
        }
    }

    fn initialize(&mut self) -> Result<()> {
        // Xavier initialization
        let fan_in = self.weights.ncols() as f64;
        let fan_out = self.weights.nrows() as f64;
        let scale = (6.0 / (fan_in + fan_out)).sqrt();

        self.weights = Array2::from_shape_fn(self.weights.raw_dim(), |_| {
            T::from(scirs2_core::random::Random::default().random::<f64>() * scale * 2.0 - scale)
                .expect("xavier initialization conversion failed")
        });

        Ok(())
    }

    fn forward(&self, input: &Array1<T>) -> Result<Array1<T>> {
        let linear_output = self.weights.dot(input) + &self.bias;
        Ok(self.apply_activation(&linear_output))
    }

    fn apply_activation(&self, x: &Array1<T>) -> Array1<T> {
        match self.activation {
            ActivationFunction::ReLU => x.mapv(|xi| if xi > T::zero() { xi } else { T::zero() }),
            ActivationFunction::GELU => x.mapv(|xi| {
                let x_f64 = xi.to_f64().unwrap_or(0.0);
                let gelu_val = 0.5
                    * x_f64
                    * (1.0 + (x_f64 * 0.7978845608 * (1.0 + 0.044715 * x_f64 * x_f64)).tanh());
                scirs2_core::numeric::NumCast::from(gelu_val).unwrap_or_else(|| T::zero())
            }),
            ActivationFunction::Swish => x.mapv(|xi| {
                let sigmoid = T::one() / (T::one() + (-xi).exp());
                xi * sigmoid
            }),
            ActivationFunction::Tanh => x.mapv(|xi| xi.tanh()),
            ActivationFunction::Sigmoid => x.mapv(|xi| T::one() / (T::one() + (-xi).exp())),
        }
    }
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> ArchitectureEncoder<T> {
    fn new(embeddingdim: usize) -> Self {
        Self {
            encoding_weights: Array2::zeros((embeddingdim, 64)), // Assume max 64 components
            _embeddingdim: embeddingdim,
            max_components: 64,
        }
    }

    fn encode(&self, architecture: &OptimizerArchitecture<T>) -> Result<Array1<T>> {
        // Simple encoding: one-hot component types + hyperparameters
        let mut encoding = Vec::new();

        for (i, component) in architecture.components.iter().enumerate() {
            if i >= self.max_components {
                break;
            }

            // Encode component type (use string hash as placeholder)
            let hash = component.bytes().fold(0u8, |acc, b| acc.wrapping_add(b));
            encoding.push(scirs2_core::numeric::NumCast::from(hash).unwrap_or_else(|| T::zero()));
        }

        // Pad to fixed size
        encoding.resize(self._embeddingdim, T::zero());
        Ok(Array1::from_vec(encoding))
    }
}

impl<T: Float + Debug + Send + Sync + 'static + Default + Clone> SearchOptimizer<T> {
    fn new(optimizer_type: SearchOptimizerType, learningrate: T) -> Self {
        Self {
            optimizer_type,
            _learningrate: learningrate,
            momentum: scirs2_core::numeric::NumCast::from(0.9).unwrap_or_else(|| T::zero()),
            parameters: HashMap::new(),
        }
    }
}