optirs-learned 0.3.2

//! Meta-SGD meta-learning algorithm
//!
//! Implements the Meta-SGD algorithm which learns per-parameter learning rates
//! in addition to the initial parameters. This allows the model to adapt
//! different parameters at different rates during task adaptation.

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

use super::functions::MetaLearner;
use super::linear_model;
use super::metrics::BatchObservations;
use super::types::{
    AdaptationStep, MetaLearningAlgorithm, MetaTask, MetaTrainingResult, QueryEvaluationMetrics,
    QueryEvaluationResult, TaskAdaptationMetrics, TaskAdaptationResult, TaskType,
};

/// Result type for inner loop adaptation: adapted parameters and adaptation trajectory
type InnerLoopResult<T> = (HashMap<String, Array1<T>>, Vec<AdaptationStep<T>>);

/// Meta-SGD learner with learnable per-parameter learning rates
///
/// Meta-SGD extends MAML by learning not just the initial parameters but also
/// per-parameter learning rates, enabling faster and more effective adaptation.
pub struct MetaSGDLearner<T: Float + Debug + Send + Sync + 'static> {
    /// Base learning rate (used to initialize per-param LRs)
    base_lr: T,
    /// Learning rate for updating the per-parameter LRs
    alpha_lr: T,
    /// Outer-loop interpolation rate toward the adapted parameters
    outer_lr: T,
    /// Number of inner loop steps
    inner_steps: usize,
    /// Learnable per-parameter learning rates
    per_param_lr: HashMap<String, Array1<T>>,
    /// Step count for tracking
    step_count: usize,
}

impl<T: Float + Debug + Send + Sync + 'static> MetaSGDLearner<T> {
    /// Create a new MetaSGDLearner with the given base learning rate
    pub fn new(base_lr: T) -> Self {
        Self {
            base_lr,
            alpha_lr: T::from(0.001).unwrap_or_else(|| T::zero()),
            outer_lr: T::from(0.1).unwrap_or_else(|| T::zero()),
            inner_steps: 5,
            per_param_lr: HashMap::new(),
            step_count: 0,
        }
    }

    /// Set the alpha learning rate for updating per-param LRs (builder pattern)
    pub fn with_alpha_lr(mut self, lr: T) -> Self {
        self.alpha_lr = lr;
        self
    }

    /// Set the outer-loop rate used to move the meta-parameters toward the
    /// adapted parameters (builder pattern).
    pub fn with_outer_lr(mut self, lr: T) -> Self {
        self.outer_lr = lr;
        self
    }

    /// Set the number of inner loop steps (builder pattern)
    pub fn with_inner_steps(mut self, n: usize) -> Self {
        self.inner_steps = n;
        self
    }

    /// Initialize per-parameter learning rates if not already set
    fn ensure_per_param_lr(&mut self, parameters: &HashMap<String, Array1<T>>) {
        for (name, param) in parameters {
            if !self.per_param_lr.contains_key(name) {
                let lr_array = Array1::from_elem(param.len(), self.base_lr);
                self.per_param_lr.insert(name.clone(), lr_array);
            }
        }
    }

    /// Compute MSE loss on a dataset given parameters
    fn compute_loss(
        &self,
        features: &[Array1<T>],
        targets: &[T],
        parameters: &HashMap<String, Array1<T>>,
    ) -> Result<T> {
        if features.is_empty() {
            return Ok(T::zero());
        }
        let mut total_loss = T::zero();
        for (feat, target) in features.iter().zip(targets.iter()) {
            let prediction = self.predict_single(feat, parameters)?;
            let diff = prediction - *target;
            total_loss = total_loss + diff * diff;
        }
        let n = T::from(features.len()).ok_or_else(|| {
            OptimError::ComputationError("Failed to convert dataset size".to_string())
        })?;
        Ok(total_loss / n)
    }

    /// Make a single prediction: weighted sum of features using parameters
    fn predict_single(
        &self,
        features: &Array1<T>,
        parameters: &HashMap<String, Array1<T>>,
    ) -> Result<T> {
        let feat_len = T::from(features.len()).ok_or_else(|| {
            OptimError::ComputationError("Failed to convert feature length".to_string())
        })?;
        if let Some(weights) = parameters.get("weights") {
            let min_len = features.len().min(weights.len());
            let mut sum = T::zero();
            for i in 0..min_len {
                sum = sum + features[i] * weights[i];
            }
            Ok(sum / feat_len)
        } else {
            let sum: T = features.iter().copied().fold(T::zero(), |a, b| a + b);
            Ok(sum / feat_len)
        }
    }

    /// Compute finite-difference gradients of loss w.r.t. parameters
    fn compute_gradients(
        &self,
        parameters: &HashMap<String, Array1<T>>,
        features: &[Array1<T>],
        targets: &[T],
    ) -> Result<HashMap<String, Array1<T>>> {
        let epsilon = T::from(1e-5)
            .ok_or_else(|| OptimError::ComputationError("Failed to convert epsilon".to_string()))?;
        let two = T::from(2.0)
            .ok_or_else(|| OptimError::ComputationError("Failed to convert 2.0".to_string()))?;
        let mut gradients = HashMap::new();

        for (name, param) in parameters {
            let mut grad = Array1::zeros(param.len());
            for i in 0..param.len() {
                let mut params_plus = parameters.clone();
                let p_plus = params_plus.get_mut(name).ok_or_else(|| {
                    OptimError::ComputationError(format!("Parameter {} not found", name))
                })?;
                p_plus[i] = p_plus[i] + epsilon;

                let mut params_minus = parameters.clone();
                let p_minus = params_minus.get_mut(name).ok_or_else(|| {
                    OptimError::ComputationError(format!("Parameter {} not found", name))
                })?;
                p_minus[i] = p_minus[i] - epsilon;

                let loss_plus = self.compute_loss(features, targets, &params_plus)?;
                let loss_minus = self.compute_loss(features, targets, &params_minus)?;

                grad[i] = (loss_plus - loss_minus) / (two * epsilon);
            }
            gradients.insert(name.clone(), grad);
        }
        Ok(gradients)
    }

    /// Run inner loop with per-parameter learning rates
    fn run_inner_loop(
        &self,
        task: &MetaTask<T>,
        initial_params: &HashMap<String, Array1<T>>,
        num_steps: usize,
    ) -> Result<InnerLoopResult<T>> {
        let mut params = initial_params.clone();
        let mut trajectory = Vec::new();

        for step in 0..num_steps {
            let loss = self.compute_loss(
                &task.support_set.features,
                &task.support_set.targets,
                &params,
            )?;
            let gradients = self.compute_gradients(
                &params,
                &task.support_set.features,
                &task.support_set.targets,
            )?;

            let grad_norm = gradients
                .values()
                .flat_map(|g| g.iter().copied())
                .map(|v| v * v)
                .fold(T::zero(), |a, b| a + b);

            // Update using per-parameter learning rates: params -= per_param_lr * grad
            let mut param_change_sq = T::zero();
            for (name, param) in params.iter_mut() {
                if let Some(grad) = gradients.get(name) {
                    let lr = self.per_param_lr.get(name);
                    for i in 0..param.len() {
                        let effective_lr = lr
                            .and_then(|lr_arr| {
                                if i < lr_arr.len() {
                                    Some(lr_arr[i])
                                } else {
                                    None
                                }
                            })
                            .unwrap_or(self.base_lr);
                        let change = effective_lr * grad[i];
                        param[i] = param[i] - change;
                        param_change_sq = param_change_sq + change * change;
                    }
                }
            }

            trajectory.push(AdaptationStep {
                step,
                loss,
                gradient_norm: grad_norm,
                parameter_change_norm: param_change_sq,
                learning_rate: self.base_lr,
            });
        }

        Ok((params, trajectory))
    }

    /// Clamp per-parameter learning rates to a valid range
    fn clamp_per_param_lr(&mut self) {
        let min_lr = T::from(1e-6).unwrap_or_else(|| T::zero());
        let max_lr = T::from(1.0).unwrap_or_else(|| T::one());
        for lr_arr in self.per_param_lr.values_mut() {
            for i in 0..lr_arr.len() {
                if lr_arr[i] < min_lr {
                    lr_arr[i] = min_lr;
                }
                if lr_arr[i] > max_lr {
                    lr_arr[i] = max_lr;
                }
            }
        }
    }
}

impl<
        T: Float
            + Debug
            + Send
            + Sync
            + 'static
            + Default
            + Clone
            + std::iter::Sum
            + scirs2_core::ndarray::ScalarOperand,
    > MetaLearner<T> for MetaSGDLearner<T>
{
    fn meta_train_step(
        &mut self,
        task_batch: &[MetaTask<T>],
        meta_parameters: &mut HashMap<String, Array1<T>>,
    ) -> Result<MetaTrainingResult<T>> {
        if task_batch.is_empty() {
            return Err(OptimError::InsufficientData("Empty task batch".to_string()));
        }

        // Initialize per-param LRs if needed
        self.ensure_per_param_lr(meta_parameters);

        let batch_size = T::from(task_batch.len()).ok_or_else(|| {
            OptimError::ComputationError("Failed to convert batch size".to_string())
        })?;

        let mut total_meta_loss = T::zero();
        let mut task_losses = Vec::new();
        let mut meta_gradients: HashMap<String, Array1<T>> = HashMap::new();

        // Accumulated parameter updates
        let mut accumulated_diff: HashMap<String, Array1<T>> = HashMap::new();
        for (name, param) in meta_parameters.iter() {
            accumulated_diff.insert(name.clone(), Array1::zeros(param.len()));
        }

        // Real per-task observations feed the reported metrics (previously a
        // block of hard-coded constants regardless of what actually happened).
        let mut observations = BatchObservations::<T>::new();

        for task in task_batch {
            // Save initial params for computing change
            let initial_params = meta_parameters.clone();

            // Pre-adaptation query loss, before the inner loop runs.
            let pre_loss = self.compute_loss(
                &task.query_set.features,
                &task.query_set.targets,
                &initial_params,
            )?;

            // Run inner loop with per-param LRs
            let (adapted_params, _trajectory) =
                self.run_inner_loop(task, meta_parameters, self.inner_steps)?;

            // Compute task loss on query set
            let task_loss = self.compute_loss(
                &task.query_set.features,
                &task.query_set.targets,
                &adapted_params,
            )?;
            task_losses.push(task_loss);
            total_meta_loss = total_meta_loss + task_loss;

            // Compute gradients on query set with adapted params
            let query_gradients = self.compute_gradients(
                &adapted_params,
                &task.query_set.features,
                &task.query_set.targets,
            )?;

            // Update per_param_lr: per_param_lr -= alpha_lr * grad * param_change
            for (name, adapted_param) in &adapted_params {
                if let Some(initial_param) = initial_params.get(name) {
                    if let Some(query_grad) = query_gradients.get(name) {
                        if let Some(lr_arr) = self.per_param_lr.get_mut(name) {
                            for i in 0..lr_arr.len().min(adapted_param.len()) {
                                let param_change = adapted_param[i] - initial_param[i];
                                lr_arr[i] =
                                    lr_arr[i] - self.alpha_lr * query_grad[i] * param_change;
                            }
                        }
                    }
                }
            }

            // Accumulate difference for meta-parameter update
            for (name, adapted_param) in &adapted_params {
                if let Some(meta_param) = meta_parameters.get(name) {
                    let diff = adapted_param - meta_param;
                    if let Some(acc) = accumulated_diff.get_mut(name) {
                        *acc = acc.clone() + &diff;
                    }
                }
            }

            // Record this task's observations for the batch metrics.
            let mut param_travel = T::zero();
            let mut task_grad: HashMap<String, Array1<T>> = HashMap::new();
            for (name, adapted_param) in &adapted_params {
                if let Some(initial_param) = initial_params.get(name) {
                    let diff = adapted_param - initial_param;
                    param_travel =
                        param_travel + diff.iter().fold(T::zero(), |a, &v| a + v * v).sqrt();
                    // Descent-convention per-task meta-gradient: meta - adapted.
                    task_grad.insert(name.clone(), diff.mapv(|v| -v));
                }
            }
            let grad_norm = query_gradients
                .values()
                .flat_map(|g| g.iter())
                .fold(T::zero(), |a, &v| a + v * v)
                .sqrt();

            observations.pre_losses.push(pre_loss);
            observations.post_losses.push(task_loss);
            observations.convergence_steps.push(self.inner_steps);
            observations.parameter_changes.push(param_travel);
            observations.gradient_norms.push(grad_norm);
            observations.gradients.push(task_grad);
        }

        // Clamp per-param LRs to valid range
        self.clamp_per_param_lr();

        // Update meta-parameters toward the adapted parameters.
        //
        // Update contract: the learner owns this update, and `meta_gradients`
        // is reported in *descent* convention (the vector `g` for which the
        // update is `theta <- theta - outer_lr * g`). Callers must not apply it
        // a second time.
        let outer_lr = self.outer_lr;
        for (name, param) in meta_parameters.iter_mut() {
            if let Some(acc) = accumulated_diff.get(name) {
                let avg_diff = acc / batch_size;
                for i in 0..param.len() {
                    param[i] = param[i] + outer_lr * avg_diff[i];
                }
                meta_gradients.insert(name.clone(), avg_diff.mapv(|v| -v));
            }
        }

        let meta_loss = total_meta_loss / batch_size;
        self.step_count += 1;

        Ok(MetaTrainingResult {
            meta_loss,
            task_losses,
            meta_gradients,
            metrics: observations.training_metrics(),
            adaptation_stats: observations.adaptation_statistics(),
        })
    }

    fn adapt_to_task(
        &mut self,
        task: &MetaTask<T>,
        meta_parameters: &HashMap<String, Array1<T>>,
        adaptation_steps: usize,
    ) -> Result<TaskAdaptationResult<T>> {
        self.ensure_per_param_lr(meta_parameters);

        let (adapted_parameters, adaptation_trajectory) =
            self.run_inner_loop(task, meta_parameters, adaptation_steps)?;

        let initial_loss = adaptation_trajectory
            .first()
            .map(|s| s.loss)
            .unwrap_or_else(T::zero);
        let final_loss = self.compute_loss(
            &task.support_set.features,
            &task.support_set.targets,
            &adapted_parameters,
        )?;

        // Every field below is measured from the trajectory this call just
        // produced, matching what `ReptileLearner`/`MAMLLearner` report.
        let steps_t = T::from(adaptation_steps.max(1)).ok_or_else(|| {
            OptimError::ComputationError("Failed to convert adaptation steps".to_string())
        })?;
        let improvement = initial_loss - final_loss;
        let travel = adaptation_trajectory
            .iter()
            .map(|s| s.parameter_change_norm)
            .fold(T::zero(), |a, b| a + b);
        let mut non_increasing = 0usize;
        for pair in adaptation_trajectory.windows(2) {
            if pair[1].loss <= pair[0].loss {
                non_increasing += 1;
            }
        }

        Ok(TaskAdaptationResult {
            adapted_parameters,
            metrics: TaskAdaptationMetrics {
                // Mean loss reduction per inner step.
                convergence_speed: improvement / steps_t,
                // Fraction of the initial loss that adaptation removed.
                final_performance: if initial_loss > T::zero() {
                    (improvement / initial_loss).max(T::zero()).min(T::one())
                } else {
                    T::zero()
                },
                // Loss reduction bought per unit of parameter travel.
                efficiency: if travel > T::zero() {
                    improvement / travel
                } else {
                    T::zero()
                },
                // Fraction of steps that did not make the loss worse.
                robustness: if adaptation_trajectory.len() > 1 {
                    let denom = T::from(adaptation_trajectory.len() - 1).unwrap_or_else(T::one);
                    T::from(non_increasing).unwrap_or_else(T::zero) / denom
                } else {
                    T::one()
                },
            },
            adaptation_trajectory,
            final_loss,
        })
    }

    fn evaluate_query_set(
        &self,
        task: &MetaTask<T>,
        adapted_parameters: &HashMap<String, Array1<T>>,
    ) -> Result<QueryEvaluationResult<T>> {
        if task.query_set.features.is_empty() {
            return Err(OptimError::InsufficientData(format!(
                "task '{}' has an empty query set",
                task.id
            )));
        }

        let mut predictions = Vec::with_capacity(task.query_set.features.len());
        let mut total_loss = T::zero();

        for (features, target) in task.query_set.features.iter().zip(&task.query_set.targets) {
            let prediction = self.predict_single(features, adapted_parameters)?;
            let diff = prediction - *target;
            total_loss = total_loss + diff * diff;
            predictions.push(prediction);
        }

        let n = T::from(task.query_set.features.len()).ok_or_else(|| {
            OptimError::ComputationError("Failed to convert query set size".to_string())
        })?;
        let query_loss = total_loss / n;

        // Confidence is the squashed per-sample residual: 1 for an exact hit.
        let confidence_scores: Vec<T> = predictions
            .iter()
            .zip(task.query_set.targets.iter())
            .map(|(p, y)| T::one() / (T::one() + (*p - *y).abs()))
            .collect();

        let classification_accuracy = match task.task_type {
            TaskType::Classification => {
                linear_model::label_accuracy(&predictions, &task.query_set.targets)
            }
            _ => None,
        };
        // Task-appropriate goodness: label accuracy for classification, R^2
        // clamped to [0, 1] for regression, 0 when neither is defined.
        let accuracy = classification_accuracy
            .or_else(|| {
                linear_model::r_squared(&predictions, &task.query_set.targets)
                    .map(|r| r.max(T::zero()).min(T::one()))
            })
            .unwrap_or_else(T::zero);

        // Full-curve AUC-ROC ranked on the retained predictions; `None` unless
        // the targets are genuinely binary with both classes represented.
        let auc = match task.task_type {
            TaskType::Classification => {
                linear_model::roc_auc(&predictions, &task.query_set.targets)
            }
            _ => None,
        };

        let count = T::from(confidence_scores.len()).ok_or_else(|| {
            OptimError::ComputationError("Failed to convert confidence count".to_string())
        })?;
        let uncertainty_quality = confidence_scores
            .iter()
            .copied()
            .fold(T::zero(), |a, b| a + b)
            / count;

        Ok(QueryEvaluationResult {
            query_loss,
            accuracy,
            predictions,
            confidence_scores,
            metrics: QueryEvaluationMetrics {
                mse: Some(query_loss),
                classification_accuracy,
                auc,
                uncertainty_quality,
            },
        })
    }

    fn get_algorithm(&self) -> MetaLearningAlgorithm {
        MetaLearningAlgorithm::MetaSGD
    }
}

#[cfg(test)]
mod tests {
    use super::super::types::{DatasetMetadata, TaskDataset, TaskMetadata, TaskType};
    use super::*;
    use scirs2_core::ndarray::Array1;

    fn make_test_task() -> MetaTask<f64> {
        let support_features = vec![
            Array1::from_vec(vec![1.0, 2.0, 3.0]),
            Array1::from_vec(vec![4.0, 5.0, 6.0]),
            Array1::from_vec(vec![7.0, 8.0, 9.0]),
        ];
        let support_targets = vec![2.0, 5.0, 8.0];
        let query_features = vec![
            Array1::from_vec(vec![2.0, 3.0, 4.0]),
            Array1::from_vec(vec![5.0, 6.0, 7.0]),
        ];
        let query_targets = vec![3.0, 6.0];

        MetaTask {
            id: "test_task".to_string(),
            support_set: TaskDataset {
                features: support_features,
                targets: support_targets,
                weights: vec![1.0, 1.0, 1.0],
                metadata: DatasetMetadata::default(),
            },
            query_set: TaskDataset {
                features: query_features,
                targets: query_targets,
                weights: vec![1.0, 1.0],
                metadata: DatasetMetadata::default(),
            },
            metadata: TaskMetadata::default(),
            difficulty: 1.0,
            domain: "test".to_string(),
            task_type: TaskType::Regression,
        }
    }

    fn make_test_params() -> HashMap<String, Array1<f64>> {
        let mut params = HashMap::new();
        params.insert("weights".to_string(), Array1::from_vec(vec![0.5, 0.5, 0.5]));
        params
    }

    #[test]
    fn test_meta_sgd_new() {
        let learner = MetaSGDLearner::new(0.01f64);
        assert_eq!(learner.inner_steps, 5);
        assert_eq!(learner.step_count, 0);
        assert!(learner.per_param_lr.is_empty());
    }

    #[test]
    fn test_meta_sgd_builder() {
        let learner = MetaSGDLearner::new(0.01f64)
            .with_alpha_lr(0.005)
            .with_inner_steps(10);
        assert_eq!(learner.inner_steps, 10);
    }

    #[test]
    fn test_meta_sgd_adapt_to_task() {
        let mut learner = MetaSGDLearner::new(0.01f64).with_inner_steps(3);
        let task = make_test_task();
        let params = make_test_params();

        let result = learner
            .adapt_to_task(&task, &params, 3)
            .expect("adapt_to_task should succeed");
        assert_eq!(result.adaptation_trajectory.len(), 3);
        assert!(result.adapted_parameters.contains_key("weights"));
        // Per-param LRs should now be initialized
        assert!(learner.per_param_lr.contains_key("weights"));
    }

    #[test]
    fn test_meta_sgd_meta_train_step() {
        let mut learner = MetaSGDLearner::new(0.01f64).with_inner_steps(3);
        let task = make_test_task();
        let mut params = make_test_params();

        let original_weights = params.get("weights").expect("weights should exist").clone();
        let result = learner
            .meta_train_step(&[task], &mut params)
            .expect("meta_train_step should succeed");

        assert_eq!(result.task_losses.len(), 1);
        let updated_weights = params.get("weights").expect("weights should exist");
        let changed = original_weights
            .iter()
            .zip(updated_weights.iter())
            .any(|(a, b)| (a - b).abs() > 1e-12);
        assert!(changed, "Meta-parameters should change after training step");
    }

    /// F14 residual regression: the reported training metrics/adaptation
    /// statistics must be computed from the run, not the old hard-coded
    /// constants (avg_speed 2.0 / generalization 0.85 / diversity 0.7 /
    /// alignment 0.9 / efficiency 0.8 / forgetting 0.1).
    #[test]
    fn test_meta_sgd_metrics_are_real_not_constant() {
        let mut learner = MetaSGDLearner::new(0.01f64).with_inner_steps(3);
        let task = make_test_task();
        let mut params = make_test_params();

        let result = learner
            .meta_train_step(&[task], &mut params)
            .expect("meta_train_step should succeed");

        // Single-task batch: nothing to disperse (diversity 0), the gradient is
        // self-aligned (1.0), and forgetting is unobservable from one batch (0).
        // The fabricated code returned 0.7 / 0.9 / 0.1 regardless of the data.
        approx::assert_abs_diff_eq!(result.metrics.task_diversity, 0.0, epsilon = 1e-12);
        approx::assert_abs_diff_eq!(result.metrics.gradient_alignment, 1.0, epsilon = 1e-12);
        approx::assert_abs_diff_eq!(
            result.adaptation_stats.stability_metrics.forgetting_measure,
            0.0,
            epsilon = 1e-12
        );
        // Statistics reflect the real run.
        assert_eq!(result.adaptation_stats.convergence_steps, vec![3]);
        assert_eq!(result.adaptation_stats.final_losses, result.task_losses);
        assert!((0.0..=1.0).contains(&result.metrics.generalization_performance));
        assert!(result.metrics.avg_adaptation_speed.is_finite());
    }

    #[test]
    fn test_meta_sgd_per_param_lr_update() {
        let mut learner = MetaSGDLearner::new(0.01f64).with_inner_steps(3);
        let task = make_test_task();
        let mut params = make_test_params();

        // First step initializes per-param LRs
        learner
            .meta_train_step(std::slice::from_ref(&task), &mut params)
            .expect("first step should succeed");

        let lr_after_first = learner
            .per_param_lr
            .get("weights")
            .expect("should have weights lr")
            .clone();

        // Second step should update per-param LRs
        let task2 = make_test_task();
        learner
            .meta_train_step(&[task2], &mut params)
            .expect("second step should succeed");

        let lr_after_second = learner
            .per_param_lr
            .get("weights")
            .expect("should have weights lr");

        // Per-param LRs should have changed between steps
        let lr_changed = lr_after_first
            .iter()
            .zip(lr_after_second.iter())
            .any(|(a, b)| (a - b).abs() > 1e-15);
        assert!(
            lr_changed,
            "Per-parameter learning rates should update during training"
        );
    }

    #[test]
    fn test_meta_sgd_evaluate_query_set() {
        let learner = MetaSGDLearner::new(0.01f64);
        let task = make_test_task();
        let params = make_test_params();

        let result = learner
            .evaluate_query_set(&task, &params)
            .expect("evaluate_query_set should succeed");
        assert_eq!(result.predictions.len(), 2);
        assert!(result.query_loss >= 0.0);
    }

    #[test]
    fn test_meta_sgd_get_algorithm() {
        let learner = MetaSGDLearner::new(0.01f64);
        assert!(matches!(
            learner.get_algorithm(),
            MetaLearningAlgorithm::MetaSGD
        ));
    }

    #[test]
    fn test_meta_sgd_empty_batch_error() {
        let mut learner = MetaSGDLearner::new(0.01f64);
        let mut params = make_test_params();
        let result = learner.meta_train_step(&[], &mut params);
        assert!(result.is_err());
    }
}