//! Complete Image Classification Example
//!
//! This example demonstrates how to build, train, and evaluate an image classifier
//! using the scirs2-neural library. It covers:
//! - Building CNN architectures with the high-level API
//! - Data loading and preprocessing
//! - Training with modern techniques (data augmentation, learning rate scheduling)
//! - Model evaluation and visualization
//! - Saving and loading trained models

use ndarray::{s, Array, Array4, ArrayD, Axis};
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use scirs2_neural::augmentation::AugmentationManager;
// use scirs2_neural::callbacks::{EarlyStopping, ModelCheckpoint};
use scirs2_neural::data::Dataset;
// use scirs2_neural::error::Result;
// use scirs2_neural::evaluation::{EvaluationConfig, Evaluator};
use scirs2_neural::prelude::*;
// ModelSerializer import removed - not available in current API
use scirs2_neural::training::{Trainer, TrainingConfig, ValidationSettings};
// use scirs2_neural::visualization::TrainingVisualizer;
// use std::path::Path;
/// Simple dataset for demonstration - CIFAR-like synthetic images
#[derive(Clone)]
struct SyntheticImageDataset {
    images: Array4<f32>,
    labels: Vec<usize>,
    num_classes: usize,
}
impl SyntheticImageDataset {
    /// Create a new synthetic dataset
    fn new(_num_samples: usize, num_classes: usize, imagesize: (usize, usize)) -> Self {
        let mut rng = SmallRng::from_seed([42; 32]);
        let channels = 3; // RGB images
        let mut images = Array4::zeros((_num_samples, channels, image_size.0, image_size.1));
        let mut labels = Vec::with_capacity(_num_samples);
        // Generate synthetic patterns for each class
        let class_patterns = (0..num_classes)
            .map(|class_id| {
                // Create unique color patterns for each class
                let r_bias = (class_id as f32 / num_classes as f32) * 0.5 + 0.25;
                let g_bias =
                    ((class_id * 3) % num_classes) as f32 / num_classes as f32 * 0.5 + 0.25;
                let b_bias =
                    ((class_id * 7) % num_classes) as f32 / num_classes as f32 * 0.5 + 0.25;
                (r_bias, g_bias, b_bias)
            })
            .collect::<Vec<_>>();
        for i in 0..num_samples {
            let class = rng.random_range(0..num_classes);
            let (r_bias..g_bias, b_bias) = class_patterns[class];
            for h in 0..image_size.0 {
                for w in 0..image_size.1 {
                    // Add spatial patterns based on position
                    let spatial_pattern = ((h + w) % 4) as f32 / 4.0;
                    // Generate pixel values with class-specific bias and spatial pattern
                    images[[i, 0, h, w]] =
                        (r_bias + spatial_pattern * 0.3 + rng.random::<f32>() * 0.2).min(1.0);
                    images[[i, 1, h, w]] =
                        (g_bias + spatial_pattern * 0.3 + rng.random::<f32>() * 0.2).min(1.0);
                    images[[i, 2, h, w]] =
                        (b_bias + spatial_pattern * 0.3 + rng.random::<f32>() * 0.2).min(1.0);
                }
            }
            labels.push(class);
        }
        Self {
            images,
            labels,
            num_classes,
    }
    /// Split dataset into train and validation sets
    fn train_val_split(&self, valratio: f32) -> (Self, Self) {
        let total_samples = self.len();
        let val_size = (total_samples as f32 * val_ratio) as usize;
        let train_size = total_samples - val_size;
        let train_images = self.images.slice(s![0..train_size, .., .., ..]).to_owned();
        let train_labels = self.labels[0..train_size].to_vec();
        let val_images = self.images.slice(s![train_size.., .., .., ..]).to_owned();
        let val_labels = self.labels[train_size..].to_vec();
        let train_dataset = Self {
            images: train_images,
            labels: train_labels,
            num_classes: self.num_classes,
        };
        let val_dataset = Self {
            images: val_images,
            labels: val_labels,
        (train_dataset, val_dataset)
impl Dataset<f32> for SyntheticImageDataset {
    fn len(&self) -> usize {
        self.images.shape()[0]
    fn get(&self, index: usize) -> scirs2_neural::error::Result<(ArrayD<f32>, ArrayD<f32>)> {
        if index >= self.len() {
            return Err(scirs2_neural::error::NeuralError::InvalidArgument(format!(
                "Index {} out of bounds for dataset of size {}",
                index,
                self.len()
            )));
        let (_, channels, height, width) = self.images.dim();
        let mut image = Array4::zeros((1, channels, height, width));
        let mut label = Array::zeros((1, self.num_classes));
        // Copy image data
        for c in 0..channels {
            for h in 0..height {
                for w in 0..width {
                    image[[0, c, h, w]] = self.images[[index, c, h, w]];
        // One-hot encode label
        label[[0, self.labels[index]]] = 1.0;
        Ok((image.into_dyn(), label.into_dyn()))
    fn box_clone(&self) -> Box<dyn Dataset<f32> + Send + Sync> {
        Box::new(self.clone())
/// Build a CNN model for image classification
#[allow(dead_code)]
fn build_cnn_model(
    input_channels: usize,
    rng: &mut SmallRng,
) -> Result<Sequential<f32>> {
    let mut model = Sequential::new();
    // First convolutional block
    model.add(Dense::new(
        input_channels * 32 * 32,
        512,
        Some("relu"),
        rng,
    )?);
    model.add(Dropout::new(0.25, rng)?);
    // Hidden layers
    model.add(Dense::new(512, 256, Some("relu"), rng)?);
    model.add(Dense::new(256, 128, Some("relu"), rng)?);
    // Output layer
    model.add(Dense::new(128, num_classes, Some("softmax"), rng)?);
    Ok(model)
/// Create training configuration with modern techniques
#[allow(dead_code)]
fn create_training_config() -> TrainingConfig {
    TrainingConfig {
        batch_size: 32,
        epochs: 50,
        learning_rate: 0.001,
        shuffle: true,
        verbose: 1,
        validation: Some(ValidationSettings {
            enabled: true,
            validation_split: 0.2,
            batch_size: 64, // Larger batch for validation (no gradients needed)
            num_workers: 0,
        }),
        gradient_accumulation: None,
        mixed_precision: None,
        num_workers: 0,
/// Calculate accuracy from predictions and targets
#[allow(dead_code)]
fn calculate_accuracy(predictions: &ArrayD<f32>, targets: &ArrayD<f32>) -> f32 {
    let batch_size = predictions.shape()[0];
    let mut correct = 0;
    for i in 0..batch_size {
        let pred_row = predictions.slice(s![i, ..]);
        let target_row = targets.slice(s![i, ..]);
        // Find argmax for prediction and target
        let pred_class = pred_row
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).expect("Operation failed"))
            .map(|(i_)| i)
            .unwrap_or(0);
        let target_class = target_row
        if pred_class == target_class {
            correct += 1;
    correct as f32 / batch_size as f32
/// Main training function
#[allow(dead_code)]
fn train_image_classifier() -> Result<()> {
    println!("🚀 Starting Image Classification Training Example");
    println!("{}", "=".repeat(60));
    // Set up reproducible random number generator
    let mut rng = SmallRng::from_seed([42; 32]);
    // Dataset parameters
    let num_samples = 1000;
    let num_classes = 5;
    let image_size = (32, 32);
    let input_channels = 3;
    println!("📊 Dataset Configuration:");
    println!("   - Samples: {}", num_samples);
    println!("   - Classes: {}", num_classes);
    println!("   - Image Size: {}x{}", image_size.0, image_size.1);
    println!("   - Channels: {}", input_channels);
    // Create synthetic dataset
    println!("\n🔄 Creating synthetic dataset...");
    let dataset = SyntheticImageDataset::new(num_samples, num_classes, image_size);
    let (train_dataset, val_dataset) = dataset.train_val_split(0.2);
    println!("   - Training samples: {}", train_dataset.len());
    println!("   - Validation samples: {}", val_dataset.len());
    // Build model
    println!("\n🏗️  Building CNN model...");
    let model = build_cnn_model(input_channels, num_classes, &mut rng)?;
    // Count parameters
    let total_params: usize = model.params().iter().map(|p| p.len()).sum();
    println!("   - Model layers: {}", model.len());
    println!("   - Total parameters: {}", total_params);
    // Create training configuration
    let config = create_training_config();
    println!("\n⚙️  Training Configuration:");
    println!("   - Batch size: {}", config.batch_size);
    println!("   - Learning rate: {}", config.learning_rate);
    println!("   - Epochs: {}", config.epochs);
    println!(
        "   - Validation split: {:.1}%",
        config.validation.as_ref().expect("Operation failed").validation_split * 100.0
    );
    // Set up training components
    let loss_fn = CrossEntropyLoss::new(1e-7);
    let optimizer = Adam::new(config.learning_rate as f32, 0.9, 0.999, 1e-8);
    // Create trainer
    let mut trainer = Trainer::new(model, optimizer, loss_fn, config);
    // Add callbacks
    trainer.add_callback(Box::new(|| {
        // Custom callback for additional logging
        println!("🔄 Epoch completed");
        Ok(())
    }));
    // Train the model
    println!("\n🏋️  Starting training...");
    println!("{}", "-".repeat(40));
    let training_session = trainer.train(&train_dataset, Some(&val_dataset))?;
    println!("\n✅ Training completed!");
    println!("   - Epochs trained: {}", training_session.epochs_trained);
        "   - Final learning rate: {:.6}",
        training_session.initial_learning_rate
    // Evaluate on validation set
    println!("\n📊 Final Evaluation:");
    let val_metrics = trainer.validate(&val_dataset)?;
    for (metric, value) in &val_metrics {
        println!("   - {}: {:.4}", metric, value);
    // Test predictions on a few samples
    println!("\n🔍 Sample Predictions:");
    let sample_indices = vec![0, 1, 2, 3, 4];
    // Manually collect batch since get_batch is not part of Dataset trait
    let mut batch_images = Vec::new();
    let mut batch_targets = Vec::new();
    for &idx in &sample_indices {
        let (img, target) = val_dataset.get(idx)?;
        batch_images.push(img);
        batch_targets.push(target);
    // Concatenate into batch arrays
    let sample_images = ndarray::concatenate(
        Axis(0),
        &batch_images.iter().map(|a| a.view()).collect::<Vec<_>>(),
    )?;
    let sample_targets = ndarray::concatenate(
        &batch_targets.iter().map(|a| a.view()).collect::<Vec<_>>(),
    let model = trainer.get_model();
    let predictions = model.forward(&sample_images)?;
    for i in 0..sample_indices.len() {
        let target_row = sample_targets.slice(s![i, ..]);
        let confidence = pred_row[pred_class];
        println!(
            "   Sample {}: Predicted={}, Actual={}, Confidence={:.3}",
            i + 1,
            pred_class,
            target_class,
            confidence
        );
    // Calculate overall accuracy
    let overall_predictions = trainer.get_model().forward(&sample_images)?;
    let accuracy = calculate_accuracy(&overall_predictions, &sample_targets);
    println!("\n🎯 Sample Accuracy: {:.2}%", accuracy * 100.0);
    // Model summary
    println!("\n📋 Training Summary:");
    let session = trainer.get_session();
    if let Some(loss_history) = session.get_metric("loss") {
        if !loss_history.is_empty() {
            println!("   - Initial loss: {:.4}", loss_history[0]);
            println!(
                "   - Final loss: {:.4}",
                loss_history[loss_history.len() - 1]
            );
    if let Some(val_loss_history) = session.get_metric("val_loss") {
        if !val_loss_history.is_empty() {
                "   - Final validation loss: {:.4}",
                val_loss_history[val_loss_history.len() - 1]
    println!("\n🎉 Image classification example completed successfully!");
    Ok(())
/// Demonstrate data augmentation techniques
#[allow(dead_code)]
fn demonstrate_augmentation() -> Result<()> {
    println!("\n🔄 Data Augmentation Demo:");
    println!("{}", "-".repeat(30));
    // Create augmentation manager
    let _aug_manager: AugmentationManager<f32> = AugmentationManager::new(Some(42));
    // Note: Augmentation API is being updated
    // For now, demonstrate basic concept
    println!("   - Augmentation manager created with seed 42");
    println!("   - Basic augmentations (rotation, flipping, etc.) would be applied here");
    // Create sample image
    let sample_image = Array4::<f32>::ones((1, 3, 32, 32));
    println!("   - Sample image shape: {:?}", sample_image.shape());
    // Note: Apply augmentation when API is stabilized
    // let augmented = aug_manager.apply(&sample_image)?;
    println!("   - Original shape: {:?}", sample_image.shape());
    println!("   - Augmentation functionality available (API being finalized)");
    println!("   ✅ Augmentation framework initialized successfully");
/// Demonstrate model saving and loading
#[allow(dead_code)]
fn demonstrate_model_persistence() -> Result<()> {
    println!("\n💾 Model Persistence Demo:");
    let mut rng = SmallRng::from_seed(123);
    // Create a simple model
    let model = build_cnn_model(3, 5, &mut rng)?;
    // Save model model.params(.iter().map(|p| p.len()).sum::<usize>()
    println!("   ✅ Model persistence simulation completed");
/// Main function
#[allow(dead_code)]
fn main() -> Result<()> {
    // Main training example
    train_image_classifier()?;
    // Additional demonstrations
    demonstrate_augmentation()?;
    demonstrate_model_persistence()?;
    println!("\n🌟 All examples completed successfully!");
    println!("🔗 Next steps:");
    println!("   - Try with real image datasets");
    println!("   - Experiment with different architectures");
    println!("   - Add more sophisticated augmentations");
    println!("   - Implement custom loss functions");
#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_synthetic_dataset_creation() {
        let dataset = SyntheticImageDataset::new(100, 5, (16, 16));
        assert_eq!(dataset.len(), 100);
        assert_eq!(dataset.num_classes, 5);
        let (train, val) = dataset.train_val_split(0.2);
        assert_eq!(train.len(), 80);
        assert_eq!(val.len(), 20);
    fn test_model_creation() -> Result<()> {
        let model = build_cnn_model(3, 10, &mut rng)?;
        assert!(!model.is_empty());
    fn test_accuracy_calculation() {
        let predictions = Array::from_shape_vec(
            (2, 3),
            vec![
                0.1, 0.8, 0.1, // Class 1
                0.7, 0.2, 0.1, // Class 0
            ],
        )
        .expect("Operation failed")
        .into_dyn();
        let targets = Array::from_shape_vec(
                0.0, 1.0, 0.0, // Class 1
                1.0, 0.0, 0.0, // Class 0
        let accuracy = calculate_accuracy(&predictions, &targets);
        assert_eq!(accuracy, 1.0); // Both predictions are correct
