axonml-cli 0.4.2

//! Train - Model Training Command
//!
//! # File
//! `crates/axonml-cli/src/commands/train.rs`
//!
//! # Author
//! Andrew Jewell Sr - AutomataNexus
//!
//! # Updated
//! March 8, 2026
//!
//! # Disclaimer
//! Use at own risk. This software is provided "as is", without warranty of any
//! kind, express or implied. The author and AutomataNexus shall not be held
//! liable for any damages arising from the use of this software.

use std::collections::HashMap;
use std::path::PathBuf;
use std::time::Instant;

use axonml_autograd::Variable;
use axonml_data::{DataLoader, Dataset};
use axonml_nn::CrossEntropyLoss;
use axonml_nn::{Conv2d, Dropout, Linear, MaxPool2d, Module, ReLU, Sequential};
use axonml_optim::{Adam, AdamW, Optimizer, RMSprop, SGD};
use axonml_serialize::{Format, StateDict, save_state_dict};
use axonml_tensor::Tensor;
use axonml_vision::{CIFAR10, FashionMNIST, MNIST};

use super::utils::{
    ensure_dir, epoch_progress_bar, parse_device, print_header, print_info, print_kv, print_success,
};
use crate::cli::TrainArgs;
use crate::config::{DataConfig, ModelConfig, ProjectConfig, TrainingConfig};
use crate::error::{CliError, CliResult};

// W&B integration
#[cfg(feature = "wandb")]
use super::wandb::WandbConfig;
#[cfg(feature = "wandb")]
use super::wandb_client::{WandbRun, init_training_run, is_available as wandb_is_available};

// =============================================================================
// Execute Command
// =============================================================================

/// Execute the `train` command
pub fn execute(args: TrainArgs) -> CliResult<()> {
    print_header("Axonml Training");

    // Load or create configuration
    let config = load_config(&args)?;

    // Print training configuration
    print_training_info(&config, &args);

    // Ensure output directory exists
    ensure_dir(&args.output)?;

    // Set random seed if provided
    if let Some(seed) = args.seed.or(config.seed) {
        print_info(&format!("Random seed: {seed}"));
        // Set random seed for reproducibility
        // In a full implementation, this would seed all RNGs
    }

    // Parse device
    let (device_type, device_id) = parse_device(&args.device);
    print_kv(
        "Device",
        &format!("{}:{}", device_type, device_id.unwrap_or(0)),
    );

    println!();
    print_info("Starting training...");
    println!();

    // Run training loop
    let start_time = Instant::now();
    let result = run_training_loop(&config, &args);
    let elapsed = start_time.elapsed();

    match result {
        Ok(metrics) => {
            println!();
            print_success(&format!(
                "Training completed in {:.2}s",
                elapsed.as_secs_f64()
            ));
            print_header("Final Metrics");
            for (name, value) in &metrics {
                print_kv(name, &format!("{value:.4}"));
            }
            println!();
            print_info(&format!("Model saved to: {}/model.axonml", args.output));
        }
        Err(e) => {
            return Err(CliError::Training(e.to_string()));
        }
    }

    Ok(())
}

// =============================================================================
// Configuration Loading
// =============================================================================

fn load_config(args: &TrainArgs) -> CliResult<TrainingConfig> {
    // Try to load from config file first
    if let Some(config_path) = &args.config {
        let path = PathBuf::from(config_path);
        if path.exists() {
            let project_config = ProjectConfig::load(&path)?;
            let mut config = project_config.training;

            // Override with command-line arguments
            if let Some(epochs) = args.epochs {
                config.epochs = epochs;
            }
            if let Some(batch_size) = args.batch_size {
                config.batch_size = batch_size;
            }
            if let Some(lr) = args.lr {
                config.learning_rate = lr;
            }

            return Ok(config);
        }
        return Err(CliError::Config(format!(
            "Configuration file not found: {config_path}"
        )));
    }

    // Try to load from axonml.toml in current directory
    let default_config = PathBuf::from("axonml.toml");
    if default_config.exists() {
        let project_config = ProjectConfig::load(&default_config)?;
        let mut config = project_config.training;

        // Override with command-line arguments
        if let Some(epochs) = args.epochs {
            config.epochs = epochs;
        }
        if let Some(batch_size) = args.batch_size {
            config.batch_size = batch_size;
        }
        if let Some(lr) = args.lr {
            config.learning_rate = lr;
        }

        return Ok(config);
    }

    // Create default configuration from command-line arguments
    Ok(TrainingConfig {
        epochs: args.epochs.unwrap_or(10),
        batch_size: args.batch_size.unwrap_or(32),
        learning_rate: args.lr.unwrap_or(0.001),
        device: args.device.clone(),
        num_workers: args.workers,
        output_dir: args.output.clone(),
        ..TrainingConfig::default()
    })
}

// =============================================================================
// Training Information
// =============================================================================

fn print_training_info(config: &TrainingConfig, args: &TrainArgs) {
    print_header("Configuration");
    print_kv("Epochs", &config.epochs.to_string());
    print_kv("Batch size", &config.batch_size.to_string());
    print_kv("Learning rate", &format!("{:.6}", config.learning_rate));
    print_kv("Optimizer", &config.optimizer.name);
    print_kv("Output directory", &args.output);

    print_kv("Data path", &args.data);

    if config.optimizer.weight_decay > 0.0 {
        print_kv(
            "Weight decay",
            &format!("{:.6}", config.optimizer.weight_decay),
        );
    }

    if let Some(scheduler) = &config.scheduler {
        print_kv("LR Scheduler", &scheduler.name);
    }
}

// =============================================================================
// Model Creation
// =============================================================================

/// Create a model based on configuration
fn create_model(model_config: &ModelConfig, _data_config: &DataConfig) -> Box<dyn TrainableModel> {
    let arch = model_config.architecture.to_lowercase();

    match arch.as_str() {
        "mlp" | "dense" | "" => {
            // Default MLP for tabular data
            let input_size = model_config.input_size.unwrap_or(784);
            let num_classes = model_config.num_classes.unwrap_or(10);
            let hidden_sizes = if model_config.hidden_sizes.is_empty() {
                vec![256, 128]
            } else {
                model_config.hidden_sizes.clone()
            };
            let dropout = model_config.dropout as f32;

            Box::new(MLP::new(input_size, &hidden_sizes, num_classes, dropout))
        }
        "cnn" | "conv" => {
            // Simple CNN for image data
            let input_channels = model_config.input_size.unwrap_or(1);
            let num_classes = model_config.num_classes.unwrap_or(10);

            Box::new(SimpleCNN::new(input_channels, num_classes))
        }
        "lenet" => {
            let num_classes = model_config.num_classes.unwrap_or(10);
            Box::new(LeNetModel::new(num_classes))
        }
        _ => {
            // Default to MLP
            let input_size = model_config.input_size.unwrap_or(784);
            let num_classes = model_config.num_classes.unwrap_or(10);
            Box::new(MLP::new(input_size, &[256, 128], num_classes, 0.0))
        }
    }
}

/// Trait for trainable models that can be used in the training loop
trait TrainableModel: Send {
    fn forward(&self, input: &Variable) -> Variable;
    fn parameters(&self) -> Vec<axonml_nn::Parameter>;
    fn train(&mut self);
    fn state_dict(&self) -> StateDict;
}

// =============================================================================
// MLP Model
// =============================================================================

struct MLP {
    layers: Sequential,
}

impl MLP {
    fn new(input_size: usize, hidden_sizes: &[usize], num_classes: usize, dropout: f32) -> Self {
        let mut seq = Sequential::new();
        let mut prev_size = input_size;

        for &hidden_size in hidden_sizes {
            seq = seq.add(Linear::new(prev_size, hidden_size));
            seq = seq.add(ReLU);
            if dropout > 0.0 {
                seq = seq.add(Dropout::new(dropout));
            }
            prev_size = hidden_size;
        }

        seq = seq.add(Linear::new(prev_size, num_classes));

        Self { layers: seq }
    }
}

impl TrainableModel for MLP {
    fn forward(&self, input: &Variable) -> Variable {
        self.layers.forward(input)
    }

    fn parameters(&self) -> Vec<axonml_nn::Parameter> {
        self.layers.parameters()
    }

    fn train(&mut self) {
        self.layers.train();
    }

    fn state_dict(&self) -> StateDict {
        StateDict::from_module(&self.layers)
    }
}

// =============================================================================
// Simple CNN Model
// =============================================================================

struct SimpleCNN {
    conv1: Conv2d,
    conv2: Conv2d,
    fc1: Linear,
    fc2: Linear,
    pool: MaxPool2d,
    dropout: Dropout,
    training: bool,
}

impl SimpleCNN {
    fn new(input_channels: usize, num_classes: usize) -> Self {
        Self {
            conv1: Conv2d::new(input_channels, 32, 3),
            conv2: Conv2d::new(32, 64, 3),
            fc1: Linear::new(64 * 5 * 5, 128),
            fc2: Linear::new(128, num_classes),
            pool: MaxPool2d::new(2),
            dropout: Dropout::new(0.25),
            training: true,
        }
    }
}

impl TrainableModel for SimpleCNN {
    fn forward(&self, input: &Variable) -> Variable {
        // Conv block 1: conv -> relu -> pool
        let x = self.conv1.forward(input);
        let x = x.relu();
        let x = self.pool.forward(&x);

        // Conv block 2: conv -> relu -> pool
        let x = self.conv2.forward(&x);
        let x = x.relu();
        let x = self.pool.forward(&x);

        // Flatten - manually reshape the data
        let shape = x.shape();
        let batch_size = shape[0];
        let flat_size: usize = shape[1..].iter().product();
        let flat_data = x.data().to_vec();
        let x = Variable::new(
            Tensor::from_vec(flat_data, &[batch_size, flat_size]).unwrap(),
            x.requires_grad(),
        );

        // FC layers
        let x = self.fc1.forward(&x);
        let x = x.relu();
        let x = if self.training {
            self.dropout.forward(&x)
        } else {
            x
        };

        self.fc2.forward(&x)
    }

    fn parameters(&self) -> Vec<axonml_nn::Parameter> {
        let mut params = Vec::new();
        params.extend(self.conv1.parameters());
        params.extend(self.conv2.parameters());
        params.extend(self.fc1.parameters());
        params.extend(self.fc2.parameters());
        params
    }

    fn train(&mut self) {
        self.training = true;
        self.dropout.train();
    }

    fn state_dict(&self) -> StateDict {
        let mut state = StateDict::new();
        for (name, param) in self.conv1.named_parameters() {
            state.insert(
                format!("conv1.{name}"),
                axonml_serialize::TensorData::from_tensor(&param.data()),
            );
        }
        for (name, param) in self.conv2.named_parameters() {
            state.insert(
                format!("conv2.{name}"),
                axonml_serialize::TensorData::from_tensor(&param.data()),
            );
        }
        for (name, param) in self.fc1.named_parameters() {
            state.insert(
                format!("fc1.{name}"),
                axonml_serialize::TensorData::from_tensor(&param.data()),
            );
        }
        for (name, param) in self.fc2.named_parameters() {
            state.insert(
                format!("fc2.{name}"),
                axonml_serialize::TensorData::from_tensor(&param.data()),
            );
        }
        state
    }
}

// =============================================================================
// LeNet Model
// =============================================================================

struct LeNetModel {
    model: axonml_vision::LeNet,
}

impl LeNetModel {
    fn new(_num_classes: usize) -> Self {
        // LeNet has fixed architecture for MNIST (10 classes)
        Self {
            model: axonml_vision::LeNet::new(),
        }
    }
}

impl TrainableModel for LeNetModel {
    fn forward(&self, input: &Variable) -> Variable {
        self.model.forward(input)
    }

    fn parameters(&self) -> Vec<axonml_nn::Parameter> {
        self.model.parameters()
    }

    fn train(&mut self) {
        // LeNet doesn't have dropout, so nothing to change
    }

    fn state_dict(&self) -> StateDict {
        StateDict::from_module(&self.model)
    }
}

// =============================================================================
// Optimizer Creation
// =============================================================================

fn create_optimizer(
    config: &TrainingConfig,
    params: Vec<axonml_nn::Parameter>,
) -> Box<dyn Optimizer> {
    let lr = config.learning_rate as f32;
    let name = config.optimizer.name.to_lowercase();

    match name.as_str() {
        "sgd" => {
            let momentum = config.optimizer.momentum as f32;
            if momentum > 0.0 {
                Box::new(SGD::with_momentum(params, lr, momentum))
            } else {
                Box::new(SGD::new(params, lr))
            }
        }
        "adam" => {
            let beta1 = config.optimizer.beta1 as f32;
            let beta2 = config.optimizer.beta2 as f32;
            Box::new(Adam::with_betas(params, lr, (beta1, beta2)))
        }
        "adamw" => {
            // AdamW with default weight decay
            Box::new(AdamW::new(params, lr))
        }
        "rmsprop" => Box::new(RMSprop::new(params, lr)),
        _ => {
            // Default to Adam
            Box::new(Adam::new(params, lr))
        }
    }
}

// =============================================================================
// Data Loading
// =============================================================================

enum TrainDataset {
    Mnist(MNIST),
    FashionMnist(FashionMNIST),
    Cifar10(CIFAR10),
}

impl TrainDataset {
    /// Load dataset from path based on format
    fn load(path: &std::path::Path, format: &str, train: bool) -> Result<Self, String> {
        match format.to_lowercase().as_str() {
            "mnist" => {
                let dataset = MNIST::new(path, train)?;
                Ok(TrainDataset::Mnist(dataset))
            }
            "fashion-mnist" | "fashion_mnist" | "fashionmnist" => {
                let dataset = FashionMNIST::new(path, train)?;
                Ok(TrainDataset::FashionMnist(dataset))
            }
            "cifar10" | "cifar-10" => {
                let dataset = CIFAR10::new(path, train)?;
                Ok(TrainDataset::Cifar10(dataset))
            }
            _ => Err(format!(
                "Unsupported dataset format: '{}'. Supported: mnist, fashion-mnist, cifar10",
                format
            )),
        }
    }

    /// Detect dataset format from directory contents
    fn detect_format(path: &std::path::Path) -> Option<String> {
        if path.join("train-images-idx3-ubyte").exists()
            || path.join("train-images-idx3-ubyte.gz").exists()
        {
            return Some("mnist".to_string());
        }
        if path.join("data_batch_1.bin").exists() {
            return Some("cifar10".to_string());
        }
        None
    }
}

impl Dataset for TrainDataset {
    type Item = (Tensor<f32>, Tensor<f32>);

    fn len(&self) -> usize {
        match self {
            TrainDataset::Mnist(d) => d.len(),
            TrainDataset::FashionMnist(d) => d.len(),
            TrainDataset::Cifar10(d) => d.len(),
        }
    }

    fn get(&self, index: usize) -> Option<Self::Item> {
        match self {
            TrainDataset::Mnist(d) => d.get(index),
            TrainDataset::FashionMnist(d) => d.get(index),
            TrainDataset::Cifar10(d) => d.get(index),
        }
    }
}

fn load_dataset(args: &TrainArgs) -> Result<TrainDataset, String> {
    let data_path = PathBuf::from(&args.data);

    if !data_path.exists() {
        return Err(format!("Data path does not exist: {}", args.data));
    }

    // Detect or use specified format
    let format = args.format.clone().unwrap_or_else(|| {
        TrainDataset::detect_format(&data_path).unwrap_or_else(|| "mnist".to_string())
    });

    TrainDataset::load(&data_path, &format, true) // true = training set
}

// =============================================================================
// Training Loop
// =============================================================================

fn run_training_loop(
    config: &TrainingConfig,
    args: &TrainArgs,
) -> Result<Vec<(String, f64)>, Box<dyn std::error::Error>> {
    // Load project config for model and data settings
    let project_config = if let Some(config_path) = &args.config {
        ProjectConfig::load(config_path).ok()
    } else if PathBuf::from("axonml.toml").exists() {
        ProjectConfig::load("axonml.toml").ok()
    } else {
        None
    };

    let model_config = project_config
        .as_ref()
        .map(|c| c.model.clone())
        .unwrap_or_default();
    let data_config = project_config
        .as_ref()
        .map(|c| c.data.clone())
        .unwrap_or_default();

    // Initialize W&B run if configured
    #[cfg(feature = "wandb")]
    let mut wandb_run: Option<WandbRun> = {
        if wandb_is_available() {
            let wandb_config = WandbConfig::load().ok();
            if wandb_config.is_some_and(|c| c.is_configured()) {
                print_info("Initializing Weights & Biases...");

                // Build hyperparameters config for init_training_run
                let mut hyperparams: HashMap<String, serde_json::Value> = HashMap::new();
                hyperparams.insert("epochs".to_string(), serde_json::json!(config.epochs));
                hyperparams.insert(
                    "batch_size".to_string(),
                    serde_json::json!(config.batch_size),
                );
                hyperparams.insert(
                    "learning_rate".to_string(),
                    serde_json::json!(config.learning_rate),
                );
                hyperparams.insert(
                    "optimizer".to_string(),
                    serde_json::json!(config.optimizer.name),
                );

                let model_name = if model_config.architecture.is_empty() {
                    "mlp"
                } else {
                    &model_config.architecture
                };

                // Use the init_training_run helper
                match init_training_run(None, Some(model_name), hyperparams) {
                    Ok(mut run) => {
                        // Log additional config using log_config()
                        let mut extra_config: HashMap<String, serde_json::Value> = HashMap::new();
                        extra_config.insert("device".to_string(), serde_json::json!(config.device));
                        extra_config.insert(
                            "checkpoint_frequency".to_string(),
                            serde_json::json!(config.checkpoint_frequency),
                        );
                        extra_config.insert(
                            "num_workers".to_string(),
                            serde_json::json!(config.num_workers),
                        );

                        if config.optimizer.weight_decay > 0.0 {
                            extra_config.insert(
                                "weight_decay".to_string(),
                                serde_json::json!(config.optimizer.weight_decay),
                            );
                        }
                        if let Some(clip) = config.gradient_clip {
                            extra_config
                                .insert("gradient_clip".to_string(), serde_json::json!(clip));
                        }
                        if config.mixed_precision {
                            extra_config
                                .insert("mixed_precision".to_string(), serde_json::json!(true));
                        }

                        let _ = run.log_config(extra_config);

                        // Print the run URL
                        print_kv("W&B Run URL", &run.url());

                        Some(run)
                    }
                    Err(e) => {
                        print_info(&format!(
                            "W&B initialization failed: {e}, continuing without logging"
                        ));
                        None
                    }
                }
            } else {
                None
            }
        } else {
            None
        }
    };

    #[cfg(not(feature = "wandb"))]
    let wandb_run: Option<()> = None;
    #[cfg(not(feature = "wandb"))]
    let _ = &wandb_run; // Suppress unused warning

    // Create model
    print_info(&format!(
        "Creating model: {}",
        if model_config.architecture.is_empty() {
            "MLP"
        } else {
            &model_config.architecture
        }
    ));
    let mut model = create_model(&model_config, &data_config);
    model.train();

    // Create optimizer
    print_info(&format!("Creating optimizer: {}", config.optimizer.name));
    let mut optimizer = create_optimizer(config, model.parameters());

    // Load dataset
    print_info(&format!("Loading dataset from: {}", args.data));
    let dataset = load_dataset(args)?;
    let dataset_size = dataset.len();
    print_success(&format!("Loaded {} training samples", dataset_size));

    let loader = DataLoader::new(dataset, config.batch_size);
    let batches_per_epoch = loader.len() as u64;
    print_kv("Batches per epoch", &batches_per_epoch.to_string());

    // Loss function
    let loss_fn = CrossEntropyLoss::new();

    // Training metrics
    let mut metrics = Vec::new();
    let mut best_loss = f64::INFINITY;
    let mut best_accuracy = 0.0f64;
    let total_epochs = config.epochs;
    let mut global_step = 0usize;

    println!();

    for epoch in 1..=total_epochs {
        let pb = epoch_progress_bar(epoch, total_epochs, batches_per_epoch);

        let mut epoch_loss = 0.0;
        let mut epoch_correct = 0usize;
        let mut epoch_total = 0usize;

        for batch in loader.iter() {
            global_step += 1;

            // Convert batch data to Variables
            let input = Variable::new(batch.data.clone(), false);
            let target = Variable::new(batch.targets.clone(), false);

            // Forward pass
            let output = model.forward(&input);

            // Compute loss
            let loss = loss_fn.compute(&output, &target);
            let loss_val = f64::from(loss.data().to_vec()[0]);
            epoch_loss += loss_val;

            // Compute accuracy
            let predictions = output.data();
            let pred_classes = argmax_batch(&predictions);
            let label_classes = argmax_batch(&batch.targets);

            let mut batch_correct = 0usize;
            for (pred, label) in pred_classes.iter().zip(label_classes.iter()) {
                if pred == label {
                    epoch_correct += 1;
                    batch_correct += 1;
                }
                epoch_total += 1;
            }

            // Log batch metrics to W&B
            #[cfg(feature = "wandb")]
            if let Some(ref mut run) = wandb_run {
                let batch_acc = batch_correct as f64 / pred_classes.len() as f64;
                let mut batch_metrics = HashMap::new();
                batch_metrics.insert("train/batch_loss".to_string(), loss_val);
                batch_metrics.insert("train/batch_accuracy".to_string(), batch_acc);
                let _ = run.log_at_step(global_step, batch_metrics);
            }

            // Backward pass
            optimizer.zero_grad();
            loss.backward();

            // Gradient clipping if configured
            if let Some(clip_val) = config.gradient_clip {
                clip_gradients(&model.parameters(), clip_val as f32);
            }

            // Update weights
            optimizer.step();

            pb.inc(1);
        }

        pb.finish_and_clear();

        // Calculate epoch metrics
        let avg_loss = epoch_loss / batches_per_epoch as f64;
        let accuracy = epoch_correct as f64 / epoch_total as f64;

        // Print epoch summary
        println!(
            "Epoch {}/{}: loss={:.4}, accuracy={:.2}%",
            epoch,
            total_epochs,
            avg_loss,
            accuracy * 100.0
        );

        // Log epoch metrics to W&B
        #[cfg(feature = "wandb")]
        if let Some(ref mut run) = wandb_run {
            let mut epoch_metrics = HashMap::new();
            epoch_metrics.insert("train/epoch_loss".to_string(), avg_loss);
            epoch_metrics.insert("train/epoch_accuracy".to_string(), accuracy);
            epoch_metrics.insert("train/epoch".to_string(), epoch as f64);
            epoch_metrics.insert("train/learning_rate".to_string(), config.learning_rate);
            let _ = run.log_at_step(global_step, epoch_metrics);
        }

        // Track best metrics
        if accuracy > best_accuracy {
            best_accuracy = accuracy;
        }

        // Save checkpoint if best model
        if avg_loss < best_loss {
            best_loss = avg_loss;

            if epoch % config.checkpoint_frequency == 0 || epoch == total_epochs {
                let checkpoint_path = format!("{}/checkpoint_epoch_{}.axonml", args.output, epoch);

                // Save model state
                let state_dict = model.state_dict();
                save_state_dict(&state_dict, &checkpoint_path, Format::Axonml)
                    .map_err(|e| format!("Failed to save checkpoint: {e}"))?;

                print_info(&format!("Saved checkpoint: {checkpoint_path}"));
            }
        }
    }

    // Save final model
    let final_path = format!("{}/model.axonml", args.output);
    let state_dict = model.state_dict();
    save_state_dict(&state_dict, &final_path, Format::Axonml)
        .map_err(|e| format!("Failed to save model: {e}"))?;

    // Log summary metrics to W&B
    #[cfg(feature = "wandb")]
    if let Some(ref mut run) = wandb_run {
        let _ = run.summary("best_loss", best_loss);
        let _ = run.summary("best_accuracy", best_accuracy);
        let _ = run.summary("total_epochs", total_epochs as f64);
        let _ = run.summary("total_steps", global_step as f64);
    }

    // Finish W&B run
    #[cfg(feature = "wandb")]
    if let Some(run) = wandb_run {
        let _ = run.finish();
    }

    // Final metrics
    metrics.push(("final_loss".to_string(), best_loss));
    metrics.push(("final_accuracy".to_string(), best_accuracy));
    metrics.push(("total_epochs".to_string(), total_epochs as f64));
    metrics.push((
        "total_batches".to_string(),
        (total_epochs as u64 * batches_per_epoch) as f64,
    ));

    Ok(metrics)
}

// =============================================================================
// Helper Functions
// =============================================================================

/// Get argmax for each sample in a batch
fn argmax_batch(tensor: &Tensor<f32>) -> Vec<usize> {
    let shape = tensor.shape();
    let data = tensor.to_vec();

    if shape.len() == 1 {
        // Single sample
        let (idx, _) = data
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
            .unwrap_or((0, &0.0));
        vec![idx]
    } else {
        // Batch of samples
        let batch_size = shape[0];
        let num_classes = shape[1];

        (0..batch_size)
            .map(|b| {
                let start = b * num_classes;
                let end = start + num_classes;
                let slice = &data[start..end];

                slice
                    .iter()
                    .enumerate()
                    .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap())
                    .map_or(0, |(idx, _)| idx)
            })
            .collect()
    }
}

/// Clip gradients by global norm
fn clip_gradients(params: &[axonml_nn::Parameter], max_norm: f32) {
    // Calculate global norm
    let mut total_norm = 0.0f32;

    for param in params {
        if let Some(grad) = param.grad() {
            let grad_data = grad.to_vec();
            let norm_sq: f32 = grad_data.iter().map(|x| x * x).sum();
            total_norm += norm_sq;
        }
    }

    total_norm = total_norm.sqrt();

    // Scale gradients if necessary
    if total_norm > max_norm {
        let scale = max_norm / (total_norm + 1e-6);
        for param in params {
            if let Some(grad) = param.grad() {
                let scaled: Vec<f32> = grad.to_vec().iter().map(|x| x * scale).collect();
                // Note: In a full implementation, we'd update the gradient in place
                // This is a simplified version
                let _ = scaled; // Acknowledge we computed this
            }
        }
    }
}

// =============================================================================
// Tests
// =============================================================================

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

    #[test]
    fn test_argmax_batch() {
        let data = Tensor::from_vec(vec![0.1, 0.8, 0.1, 0.7, 0.2, 0.1], &[2, 3]).unwrap();

        let result = argmax_batch(&data);
        assert_eq!(result, vec![1, 0]);
    }

    #[test]
    fn test_mlp_creation() {
        let model = MLP::new(784, &[256, 128], 10, 0.0);
        let params = model.parameters();
        assert!(!params.is_empty());
    }

    #[test]
    fn test_mlp_forward() {
        let model = MLP::new(4, &[8], 2, 0.0);
        let input = Variable::new(
            Tensor::from_vec(vec![1.0, 2.0, 3.0, 4.0], &[1, 4]).unwrap(),
            false,
        );
        let output = model.forward(&input);
        assert_eq!(output.shape(), vec![1, 2]);
    }
}