torsh-optim 0.1.3

//! Neural Optimizer - Research Feature
//!
//! This module implements neural network-based optimization strategies that learn
//! to optimize functions. This is based on recent research in "learning to optimize"
//! and includes implementations of:
//! - Learning to learn by gradient descent by gradient descent (Andrychowicz et al., 2016)
//! - Learned optimizers that scale and generalize (Metz et al., 2022)
//! - MetaAdam and other learned adaptive optimizers
//!
//! WARNING: This is a research feature and may not be stable or performant
//! for production use. It is intended for experimentation and research purposes.

use crate::{Optimizer, OptimizerError, OptimizerResult};
use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use torsh_core::{
    device::{CpuDevice, DeviceType},
    DType,
};
use torsh_tensor::{creation::randn, Tensor};

/// Configuration for neural optimizer
#[derive(Debug, Clone)]
pub struct NeuralOptimizerConfig {
    /// Learning rate for the meta-optimizer (optimizer of the optimizer)
    pub meta_learning_rate: f32,
    /// Hidden size for the LSTM optimizer network
    pub hidden_size: usize,
    /// Number of layers in the optimizer network
    pub num_layers: usize,
    /// Device to run the neural optimizer on
    pub device: Arc<CpuDevice>,
    /// Maximum gradient norm for clipping
    pub max_grad_norm: f32,
    /// Whether to use coordinate-wise optimization
    pub coordinate_wise: bool,
    /// History length for the neural network
    pub history_length: usize,
}

impl Default for NeuralOptimizerConfig {
    fn default() -> Self {
        Self {
            meta_learning_rate: 0.001,
            hidden_size: 20,
            num_layers: 2,
            device: Arc::new(CpuDevice::new()),
            max_grad_norm: 10.0,
            coordinate_wise: true,
            history_length: 20,
        }
    }
}

/// Simple neural network for learning optimization updates
/// This is a simplified implementation for demonstration purposes
#[derive(Debug, Clone)]
pub struct OptimizerNetwork {
    /// LSTM-like state for each parameter
    pub hidden_states: HashMap<String, Tensor>,
    /// Cell states for LSTM
    pub cell_states: HashMap<String, Tensor>,
    /// Network weights
    pub weights: NetworkWeights,
    /// Configuration
    pub config: NeuralOptimizerConfig,
}

/// Network weights for the neural optimizer
#[derive(Debug, Clone)]
pub struct NetworkWeights {
    /// Input gate weights
    pub w_input: Tensor,
    /// Forget gate weights
    pub w_forget: Tensor,
    /// Output gate weights
    pub w_output: Tensor,
    /// Cell gate weights
    pub w_cell: Tensor,
    /// Output projection weights
    pub w_output_proj: Tensor,
    /// Biases
    pub bias_input: Tensor,
    pub bias_forget: Tensor,
    pub bias_output: Tensor,
    pub bias_cell: Tensor,
    pub bias_output_proj: Tensor,
}

impl NetworkWeights {
    /// Initialize network weights randomly
    pub fn new(input_size: usize, hidden_size: usize, device: &CpuDevice) -> OptimizerResult<Self> {
        let scale = (2.0 / (input_size + hidden_size) as f32).sqrt();

        Ok(Self {
            w_input: randn::<f32>(&[input_size + hidden_size, hidden_size])?.mul_scalar(scale)?,
            w_forget: randn::<f32>(&[input_size + hidden_size, hidden_size])?.mul_scalar(scale)?,
            w_output: randn::<f32>(&[input_size + hidden_size, hidden_size])?.mul_scalar(scale)?,
            w_cell: randn::<f32>(&[input_size + hidden_size, hidden_size])?.mul_scalar(scale)?,
            w_output_proj: randn::<f32>(&[hidden_size, 1])?.mul_scalar(scale)?,
            bias_input: Tensor::zeros(&[hidden_size], DeviceType::Cpu)?,
            bias_forget: Tensor::ones(&[hidden_size], DeviceType::Cpu)?, // Initialize forget bias to 1
            bias_output: Tensor::zeros(&[hidden_size], DeviceType::Cpu)?,
            bias_cell: Tensor::zeros(&[hidden_size], DeviceType::Cpu)?,
            bias_output_proj: Tensor::zeros(&[1], DeviceType::Cpu)?,
        })
    }
}

impl OptimizerNetwork {
    /// Create a new neural optimizer network
    pub fn new(config: NeuralOptimizerConfig) -> OptimizerResult<Self> {
        let input_size = if config.coordinate_wise {
            2 // gradient and parameter value
        } else {
            config.history_length * 2 // history of gradients and parameters
        };

        let weights = NetworkWeights::new(input_size, config.hidden_size, &config.device)?;

        Ok(Self {
            hidden_states: HashMap::new(),
            cell_states: HashMap::new(),
            weights,
            config,
        })
    }

    /// Forward pass through the neural network to compute parameter update
    pub fn forward(
        &mut self,
        param_id: &str,
        gradient: &Tensor,
        parameter: &Tensor,
    ) -> OptimizerResult<Tensor> {
        let device = self.config.device.clone();

        // Prepare input: concatenate gradient and parameter information
        let input = if self.config.coordinate_wise {
            // For coordinate-wise optimization, process each element independently
            let grad_norm = gradient.norm()?.unsqueeze(0)?;
            let param_norm = parameter.norm()?.unsqueeze(0)?;
            Tensor::cat(&[&grad_norm, &param_norm], 0)?
        } else {
            // Use full gradient and parameter vectors (simplified)
            let grad_flat = gradient.flatten()?;
            let param_flat = parameter.flatten()?;
            Tensor::cat(&[&grad_flat, &param_flat], 0)?
        };

        // Get or initialize hidden and cell states
        let hidden_shape = vec![self.config.hidden_size];
        let hidden_state = self
            .hidden_states
            .entry(param_id.to_string())
            .or_insert_with(|| {
                Tensor::zeros(&hidden_shape, DeviceType::Cpu)
                    .expect("tensor creation should succeed")
            });
        let cell_state = self
            .cell_states
            .entry(param_id.to_string())
            .or_insert_with(|| {
                Tensor::zeros(&hidden_shape, DeviceType::Cpu)
                    .expect("tensor creation should succeed")
            })
            .clone();

        // LSTM-like computation
        let combined_input = Tensor::cat(&[&input, &hidden_state.clone()], 0)?;

        // Gates computation
        let input_gate = self.sigmoid(
            &combined_input
                .matmul(&self.weights.w_input)?
                .add_op(&self.weights.bias_input)?,
        )?;
        let forget_gate = self.sigmoid(
            &combined_input
                .matmul(&self.weights.w_forget)?
                .add_op(&self.weights.bias_forget)?,
        )?;
        let output_gate = self.sigmoid(
            &combined_input
                .matmul(&self.weights.w_output)?
                .add_op(&self.weights.bias_output)?,
        )?;
        let cell_gate = self.tanh(
            &combined_input
                .matmul(&self.weights.w_cell)?
                .add_op(&self.weights.bias_cell)?,
        )?;

        // Update cell state
        let new_cell_state = forget_gate
            .mul_op(&cell_state)?
            .add_op(&input_gate.mul_op(&cell_gate)?)?;

        // Update hidden state
        let new_hidden_state = output_gate.mul_op(&self.tanh(&new_cell_state)?)?;

        // Compute parameter update
        let update_magnitude = new_hidden_state
            .matmul(&self.weights.w_output_proj)?
            .add_op(&self.weights.bias_output_proj)?;

        // Apply update to parameter shape
        let update = if self.config.coordinate_wise {
            // Scale the gradient by the learned magnitude
            gradient.mul_op(&update_magnitude.broadcast_to(gradient.shape().dims())?)?
        } else {
            // For non-coordinate-wise, we need more sophisticated reshaping
            gradient.mul_scalar(update_magnitude.item()?)?
        };

        // Update states
        *self
            .hidden_states
            .get_mut(param_id)
            .expect("hidden_states should exist for param_id") = new_hidden_state;
        *self
            .cell_states
            .get_mut(param_id)
            .expect("cell_states should exist for param_id") = new_cell_state;

        Ok(update)
    }

    /// Sigmoid activation function
    fn sigmoid(&self, x: &Tensor) -> OptimizerResult<Tensor> {
        // sigmoid(x) = 1 / (1 + exp(-x))
        let neg_x = x.mul_scalar(-1.0)?;
        let exp_neg_x = neg_x.exp()?;
        let one_plus_exp = exp_neg_x.add_scalar(1.0)?;
        Ok(one_plus_exp.reciprocal()?)
    }

    /// Tanh activation function
    fn tanh(&self, x: &Tensor) -> OptimizerResult<Tensor> {
        Ok(x.tanh()?)
    }

    /// Reset network state
    pub fn reset_state(&mut self) {
        self.hidden_states.clear();
        self.cell_states.clear();
    }

    /// Get network parameters for meta-optimization
    pub fn parameters(&self) -> Vec<&Tensor> {
        vec![
            &self.weights.w_input,
            &self.weights.w_forget,
            &self.weights.w_output,
            &self.weights.w_cell,
            &self.weights.w_output_proj,
            &self.weights.bias_input,
            &self.weights.bias_forget,
            &self.weights.bias_output,
            &self.weights.bias_cell,
            &self.weights.bias_output_proj,
        ]
    }
}

/// Neural optimizer that uses a neural network to learn optimization updates
pub struct NeuralOptimizer {
    /// Neural network for computing updates
    pub network: OptimizerNetwork,
    /// Parameters being optimized
    pub parameters: Vec<Tensor>,
    /// Meta-optimizer for training the neural optimizer
    pub meta_optimizer: Option<Box<dyn Optimizer>>,
    /// Training mode flag
    pub training: bool,
    /// Step counter
    pub step_count: usize,
}

impl NeuralOptimizer {
    /// Create a new neural optimizer
    pub fn new(
        parameters: Vec<Tensor>,
        config: Option<NeuralOptimizerConfig>,
    ) -> OptimizerResult<Self> {
        let config = config.unwrap_or_default();
        let network = OptimizerNetwork::new(config)?;

        Ok(Self {
            network,
            parameters,
            meta_optimizer: None,
            training: false,
            step_count: 0,
        })
    }

    /// Create a neural optimizer with meta-learning capabilities
    pub fn with_meta_learning(
        parameters: Vec<Tensor>,
        config: Option<NeuralOptimizerConfig>,
    ) -> OptimizerResult<Self> {
        let mut optimizer = Self::new(parameters, config)?;
        optimizer.training = true;

        // Create meta-optimizer (Adam for the neural network parameters)
        let network_params = optimizer
            .network
            .parameters()
            .iter()
            .map(|p| Arc::new(RwLock::new((*p).clone())))
            .collect();

        use crate::adam::Adam;
        let meta_optimizer = Adam::new(
            network_params,
            Some(optimizer.network.config.meta_learning_rate),
            None,
            None,
            None,
            false,
        );

        optimizer.meta_optimizer = Some(Box::new(meta_optimizer));

        Ok(optimizer)
    }

    /// Set training mode for meta-learning
    pub fn train(&mut self, mode: bool) {
        self.training = mode;
    }

    /// Reset the neural optimizer state
    pub fn reset(&mut self) {
        self.network.reset_state();
        self.step_count = 0;
    }

    /// Compute loss for meta-learning (simplified objective)
    pub fn compute_meta_loss(&self, target_loss: f32, actual_loss: f32) -> f32 {
        (target_loss - actual_loss).powi(2)
    }

    /// Update the neural network parameters using meta-gradients
    pub fn meta_step(&mut self, meta_loss: f32) -> OptimizerResult<()> {
        if let Some(ref mut meta_optimizer) = self.meta_optimizer {
            // Compute gradients of meta-loss with respect to network parameters
            // This is a simplified implementation - in practice, you'd need proper backpropagation
            for param in self.network.parameters() {
                // Simplified meta-gradient (in practice, compute actual gradients)
                let meta_grad =
                    randn::<f32>(param.shape().dims())?.mul_scalar(meta_loss * 0.001)?;
                param.set_grad(Some(meta_grad));
            }

            meta_optimizer.step()?;
        }
        Ok(())
    }
}

impl Optimizer for NeuralOptimizer {
    fn step(&mut self) -> OptimizerResult<()> {
        self.step_count += 1;

        for (i, param) in self.parameters.iter_mut().enumerate() {
            if let Some(grad) = param.grad() {
                let param_id = format!("param_{}", i);

                // Compute update using neural network
                let update = self.network.forward(&param_id, &grad, param)?;

                // Apply gradient clipping
                let update_norm = update.norm()?.item()?;
                let clipped_update = if update_norm > self.network.config.max_grad_norm {
                    update.mul_scalar(self.network.config.max_grad_norm / update_norm)?
                } else {
                    update
                };

                // Apply update to parameter
                *param = param.sub(&clipped_update)?;

                // Clear gradients
                param.set_grad(None);
            }
        }

        Ok(())
    }

    fn zero_grad(&mut self) {
        for param in &mut self.parameters {
            // Neural optimizer manages gradients internally
            // This is a placeholder implementation
        }
    }

    fn get_lr(&self) -> Vec<f32> {
        // Neural optimizer doesn't have a fixed learning rate
        vec![self.network.config.meta_learning_rate]
    }

    fn set_lr(&mut self, lr: f32) {
        // Update meta-learning rate
        self.network.config.meta_learning_rate = lr;
    }

    fn state_dict(&self) -> OptimizerResult<crate::OptimizerState> {
        let mut state = crate::OptimizerState::new("NeuralOptimizer".to_string());

        // Add meta-learning rate to global state
        state.global_state.insert(
            "meta_learning_rate".to_string(),
            self.network.config.meta_learning_rate,
        );
        state
            .global_state
            .insert("step_count".to_string(), self.step_count as f32);
        state.global_state.insert(
            "hidden_size".to_string(),
            self.network.config.hidden_size as f32,
        );
        state.global_state.insert(
            "num_layers".to_string(),
            self.network.config.num_layers as f32,
        );

        // Note: Saving/loading network weights would require more sophisticated serialization

        Ok(state)
    }

    fn add_param_group(
        &mut self,
        params: Vec<std::sync::Arc<parking_lot::RwLock<Tensor>>>,
        options: std::collections::HashMap<String, f32>,
    ) {
        // Neural optimizer manages parameters differently
        // This is a placeholder implementation
    }

    fn load_state_dict(&mut self, state: crate::OptimizerState) -> OptimizerResult<()> {
        if let Some(&meta_lr) = state.global_state.get("meta_learning_rate") {
            self.network.config.meta_learning_rate = meta_lr;
        }

        if let Some(&step_count) = state.global_state.get("step_count") {
            self.step_count = step_count as usize;
        }

        Ok(())
    }
}

/// Meta-learning trainer for neural optimizers
pub struct NeuralOptimizerTrainer {
    /// The neural optimizer being trained
    pub optimizer: NeuralOptimizer,
    /// Training tasks/problems
    pub training_tasks: Vec<Box<dyn OptimizationTask>>,
    /// Validation tasks
    pub validation_tasks: Vec<Box<dyn OptimizationTask>>,
    /// Training configuration
    pub config: TrainingConfig,
}

/// Configuration for training neural optimizers
#[derive(Debug, Clone)]
pub struct TrainingConfig {
    /// Number of meta-training iterations
    pub meta_iterations: usize,
    /// Number of inner optimization steps per task
    pub inner_steps: usize,
    /// Meta-learning rate
    pub meta_lr: f32,
    /// Device for training
    pub device: Arc<CpuDevice>,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            meta_iterations: 1000,
            inner_steps: 100,
            meta_lr: 0.001,
            device: Arc::new(CpuDevice::new()),
        }
    }
}

/// Trait for optimization tasks used in meta-learning
pub trait OptimizationTask {
    /// Initialize parameters for this task
    fn initialize_parameters(&self, device: &CpuDevice) -> OptimizerResult<Vec<Tensor>>;

    /// Compute loss and gradients for given parameters
    fn compute_loss_and_gradients(
        &self,
        parameters: &[Tensor],
    ) -> OptimizerResult<(f32, Vec<Tensor>)>;

    /// Get task name/description
    fn name(&self) -> &str;
}

/// Simple quadratic task for testing neural optimizers
pub struct QuadraticTask {
    pub dimension: usize,
    pub condition_number: f32,
    pub name: String,
}

impl QuadraticTask {
    pub fn new(dimension: usize, condition_number: f32) -> Self {
        Self {
            dimension,
            condition_number,
            name: format!("Quadratic_{}D_cond{:.1}", dimension, condition_number),
        }
    }
}

impl OptimizationTask for QuadraticTask {
    fn initialize_parameters(&self, device: &CpuDevice) -> OptimizerResult<Vec<Tensor>> {
        Ok(vec![randn::<f32>(&[self.dimension])?])
    }

    fn compute_loss_and_gradients(
        &self,
        parameters: &[Tensor],
    ) -> OptimizerResult<(f32, Vec<Tensor>)> {
        let param = &parameters[0];

        // Create ill-conditioned quadratic: f(x) = 0.5 * x^T * A * x
        // where A has eigenvalues ranging from 1 to condition_number
        let mut hessian_diag = Vec::new();
        for i in 0..self.dimension {
            let eigenval =
                1.0 + (self.condition_number - 1.0) * (i as f32) / (self.dimension as f32 - 1.0);
            hessian_diag.push(eigenval);
        }

        let hessian_diag_tensor =
            Tensor::from_data(hessian_diag, param.shape().dims().to_vec(), param.device())?;

        // Loss: 0.5 * sum(hessian_diag * x^2)
        let loss = param
            .pow(2.0)?
            .mul_op(&hessian_diag_tensor)?
            .sum()?
            .mul_scalar(0.5)?
            .item()?;

        // Gradient: hessian_diag * x
        let grad = param.mul_op(&hessian_diag_tensor)?;

        Ok((loss, vec![grad]))
    }

    fn name(&self) -> &str {
        &self.name
    }
}

impl NeuralOptimizerTrainer {
    /// Create a new trainer
    pub fn new(
        optimizer: NeuralOptimizer,
        training_tasks: Vec<Box<dyn OptimizationTask>>,
        config: Option<TrainingConfig>,
    ) -> Self {
        Self {
            optimizer,
            training_tasks,
            validation_tasks: Vec::new(),
            config: config.unwrap_or_default(),
        }
    }

    /// Train the neural optimizer on the given tasks
    pub fn train(&mut self) -> OptimizerResult<Vec<f32>> {
        let mut meta_losses = Vec::new();

        for meta_iter in 0..self.config.meta_iterations {
            let mut total_meta_loss = 0.0;

            // Sample a random task
            let task_idx = meta_iter % self.training_tasks.len();
            let task = &self.training_tasks[task_idx];

            // Initialize parameters for this task
            let mut params = task.initialize_parameters(&CpuDevice::default())?;

            // Perform inner optimization steps
            let mut task_loss = 0.0;
            for _ in 0..self.config.inner_steps {
                let (loss, grads) = task.compute_loss_and_gradients(&params)?;
                task_loss = loss;

                // Set gradients
                for (param, grad) in params.iter_mut().zip(grads.iter()) {
                    param.set_grad(Some(grad.clone()));
                }

                // Update parameters using neural optimizer
                // Note: This is simplified - in practice, you'd need to properly track the computation graph
                self.optimizer.step()?;
            }

            // Compute meta-loss (how well did we optimize?)
            let target_loss = 0.0; // Ideal target
            let meta_loss = self.optimizer.compute_meta_loss(target_loss, task_loss);
            total_meta_loss += meta_loss;

            // Update neural optimizer parameters
            self.optimizer.meta_step(meta_loss)?;

            meta_losses.push(total_meta_loss);

            if meta_iter % 100 == 0 {
                println!(
                    "Meta-iteration {}: Meta-loss = {:.6}, Task loss = {:.6} (Task: {})",
                    meta_iter,
                    meta_loss,
                    task_loss,
                    task.name()
                );
            }
        }

        Ok(meta_losses)
    }
}

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

    #[test]
    fn test_neural_optimizer_config() {
        let config = NeuralOptimizerConfig::default();
        assert_eq!(config.meta_learning_rate, 0.001);
        assert_eq!(config.hidden_size, 20);
        assert_eq!(config.num_layers, 2);
        assert_eq!(config.max_grad_norm, 10.0);
        assert!(config.coordinate_wise);
        assert_eq!(config.history_length, 20);
    }

    #[test]
    fn test_quadratic_task() {
        let task = QuadraticTask::new(10, 100.0);
        assert_eq!(task.dimension, 10);
        assert_eq!(task.condition_number, 100.0);
        assert_eq!(task.name(), "Quadratic_10D_cond100.0");
    }

    #[test]
    fn test_training_config() {
        let config = TrainingConfig::default();
        assert_eq!(config.meta_iterations, 1000);
        assert_eq!(config.inner_steps, 100);
        assert_eq!(config.meta_lr, 0.001);
    }
}