optirs-learned 0.3.2

//! Common types and configurations for learned optimizers

use crate::error::{OptimError, Result};
use scirs2_core::ndarray::{Array1, Array2};
use scirs2_core::numeric::{Float, ToPrimitive};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fmt::Debug;

/// Fallibly convert a primitive numeric value (an `f64` constant, a `usize`
/// count, ...) into a generic scalar type `T`.
///
/// This replaces the `T::from(x).expect("unwrap failed")` pattern that was
/// spread across this crate: a scalar type that genuinely cannot represent `x`
/// now produces an honest [`OptimError::InvalidConfig`] instead of panicking in
/// the middle of an optimizer step.
pub(crate) fn cast_scalar<T: Float, V: ToPrimitive + Copy + Debug>(value: V) -> Result<T> {
    T::from(value).ok_or_else(|| {
        OptimError::InvalidConfig(format!(
            "{value:?} cannot be represented in this optimizer's scalar type"
        ))
    })
}

/// [`cast_scalar`] for a value that must also be strictly positive, used for
/// divisors so a zero denominator is reported rather than producing an
/// infinity or a NaN.
pub(crate) fn cast_positive<T: Float, V: ToPrimitive + Copy + Debug>(
    value: V,
    what: &str,
) -> Result<T> {
    let converted: T = cast_scalar(value)?;
    if converted > T::zero() {
        Ok(converted)
    } else {
        Err(OptimError::InsufficientData(format!(
            "{what} must be positive, got {value:?}"
        )))
    }
}

/// Base configuration for learned optimizers
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct LearnedOptimizerConfig {
    /// Learning rate
    pub learning_rate: f64,

    /// Meta-learning rate for higher-level optimization
    pub meta_learning_rate: f64,

    /// Batch size for training
    pub batch_size: usize,

    /// Maximum number of optimization steps
    pub max_steps: usize,

    /// Convergence threshold
    pub convergence_threshold: f64,

    /// Whether to use momentum
    pub use_momentum: bool,

    /// Momentum decay factor
    pub momentum_decay: f64,

    /// Weight decay for regularization
    pub weight_decay: f64,

    /// Hidden size for neural networks
    pub hidden_size: usize,

    /// Number of attention heads
    pub attention_heads: usize,

    /// Size of gradient history buffer
    pub gradient_history_size: usize,

    /// Number of input features
    pub input_features: usize,

    /// Number of output features
    pub output_features: usize,

    /// Number of layers in neural networks
    pub num_layers: usize,

    /// Dropout rate for regularization
    pub dropout_rate: f64,

    /// Whether to use attention mechanisms
    pub use_attention: bool,

    /// Random seed for reproducibility
    pub seed: Option<u64>,
}

impl Default for LearnedOptimizerConfig {
    fn default() -> Self {
        Self {
            learning_rate: 0.001,
            meta_learning_rate: 0.0001,
            batch_size: 32,
            max_steps: 1000,
            convergence_threshold: 1e-6,
            use_momentum: true,
            momentum_decay: 0.9,
            weight_decay: 1e-4,
            hidden_size: 256,
            attention_heads: 8,
            gradient_history_size: 50,
            input_features: 256,
            output_features: 256,
            num_layers: 3,
            dropout_rate: 0.1,
            use_attention: true,
            seed: None,
        }
    }
}

/// Meta-optimization strategies
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub enum MetaOptimizationStrategy {
    /// First-order approximation (FOMAML)
    #[default]
    FirstOrder,

    /// Full second-order gradients (MAML)
    SecondOrder,

    /// MAML algorithm (alias for SecondOrder)
    MAML,

    /// Reptile algorithm
    Reptile,

    /// Custom gradient-based meta-learning
    Custom {
        inner_steps: usize,
        outer_learning_rate: f64,
    },
}

/// Base optimizer state
#[derive(Debug, Clone)]
pub struct OptimizerState<T: Float + Debug + Send + Sync + 'static> {
    /// Current parameters
    pub parameters: Array1<T>,

    /// Current gradients
    pub gradients: Array1<T>,

    /// Momentum buffer
    pub momentum: Option<Array1<T>>,

    /// Hidden states for neural optimizers
    pub hidden_states: HashMap<String, Array1<T>>,

    /// Memory buffers for attention mechanisms
    pub memory_buffers: HashMap<String, Array2<T>>,

    /// Current step number
    pub step: usize,

    /// Step count (alias for step)
    pub step_count: usize,

    /// Current loss value
    pub loss: Option<T>,

    /// Learning rate schedule
    pub learning_rate: T,

    /// Additional state metadata
    pub metadata: StateMetadata,
}

/// Learning rate a freshly-created [`OptimizerState`] starts at.
pub const DEFAULT_LEARNING_RATE: f64 = 0.001;

impl<T: Float + Debug + Send + Sync + 'static> OptimizerState<T> {
    /// A zeroed state for `num_params` parameters at
    /// [`DEFAULT_LEARNING_RATE`].
    ///
    /// # Errors
    /// Returns `Err` when the default learning rate cannot be represented in
    /// `T`. This used to panic.
    pub fn new(num_params: usize) -> Result<Self> {
        Ok(Self {
            parameters: Array1::zeros(num_params),
            gradients: Array1::zeros(num_params),
            momentum: None,
            hidden_states: HashMap::new(),
            memory_buffers: HashMap::new(),
            step: 0,
            step_count: 0,
            loss: None,
            learning_rate: cast_scalar(DEFAULT_LEARNING_RATE)?,
            metadata: StateMetadata::default(),
        })
    }
}

/// State metadata
#[derive(Debug, Clone)]
pub struct StateMetadata {
    /// Task identifier
    pub task_id: Option<String>,

    /// Optimizer type used
    pub optimizer_type: Option<String>,

    /// Version information
    pub version: String,

    /// Timestamp of creation/update
    pub timestamp: std::time::SystemTime,

    /// Checksum for integrity
    pub checksum: u64,

    /// Compression level used
    pub compression_level: u8,

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

impl Default for StateMetadata {
    fn default() -> Self {
        Self {
            task_id: None,
            optimizer_type: None,
            version: "1.0".to_string(),
            timestamp: std::time::SystemTime::now(),
            checksum: 0,
            compression_level: 0,
            custom_data: HashMap::new(),
        }
    }
}

/// Neural optimizer type variants
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum NeuralOptimizerType {
    /// Transformer-based optimizer
    Transformer,

    /// LSTM-based optimizer
    LSTM,

    /// Simple MLP optimizer
    MLP,

    /// Convolutional optimizer
    CNN,
}

/// Task context for meta-learning
#[derive(Debug, Clone)]
pub struct TaskContext<T: Float + Debug + Send + Sync + 'static> {
    /// Task identifier
    pub task_id: String,

    /// Initial parameters for the task
    pub initial_parameters: Array1<T>,

    /// Task-specific data
    pub task_data: Array2<T>,

    /// Target values for the task
    pub targets: Array1<T>,

    /// Task difficulty or complexity measure
    pub difficulty: f64,
}

/// Neural optimizer metrics
#[derive(Debug, Clone)]
pub struct NeuralOptimizerMetrics {
    /// Average loss over training
    pub avg_loss: f64,

    /// Convergence rate
    pub convergence_rate: f64,

    /// Number of steps to convergence
    pub steps_to_convergence: Option<usize>,

    /// Final accuracy/performance
    pub final_performance: f64,

    /// Training time in seconds
    pub training_time: f64,
}

/// Task performance metrics
#[derive(Debug, Clone)]
pub struct TaskPerformance {
    /// Task identifier
    pub task_id: String,

    /// Performance score (higher is better)
    pub score: f64,

    /// Whether the task converged
    pub converged: bool,

    /// Number of optimization steps taken
    pub steps_taken: usize,

    /// Final loss value
    pub final_loss: f64,
}