//! FNO Training module for neural operator demo
//!
//! Provides training functionality for Fourier Neural Operators using
//! pure Rust implementation with rtx-autograd and rtx-losses.
//!
//! # Architecture
//!
//! The trainer uses a simplified FNO architecture suitable for training:
//! - Trainable parameters stored as flat vectors
//! - Forward pass through `FNO2d` model
//! - MSE loss computation
//! - `AdamW` optimizer for parameter updates
//!
//! # Example
//!
//! ```rust,ignore
//! use rtx_neural_operator_demo::training::{FnoTrainer, FnoTrainingConfig};
//!
//! let config = FnoTrainingConfig::standard();
//! let trainer = FnoTrainer::new(config);
//!
//! // Start training (runs in background)
//! trainer.start_training().await?;
//!
//! // Poll for progress
//! let progress = trainer.get_progress();
//! ```

use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Instant;

use rtx_backend::Backend;
use rtx_backend::auto_select::{BackendType, get_backend_info, select_best_backend};
use rtx_backend_cpu::CpuBackend;
use rtx_neural_operator::FNO2d;
use rtx_neural_operator::weights::FNO2dWeights;
use rtx_neural_operator_shared::config::PDEConfig;
use rtx_neural_operator_shared::ipc::{TrainingConfig, TrainingProgress, TrainingStatus};
use rtx_nn::GenericModule4D;
use rtx_tensor::GenericTensor;
use tokio::sync::broadcast;

use crate::data::{DarcyDataConfig, DarcyDataGenerator, DarcySample, TrainingBatch};

// Conditional GPU backend imports
#[cfg(feature = "cuda")]
use rtx_backend_cuda::{CudaBackend, CudaDevice};
#[cfg(feature = "metal")]
use rtx_backend_metal::{MetalBackend, MetalDevice};
#[cfg(feature = "rocm")]
use rtx_backend_rocm::{RocmBackend, RocmDevice};

/// Get a human-readable device string for a specific backend type.
fn get_device_string_for_backend(backend_type: BackendType) -> String {
    let info = get_backend_info(backend_type);

    match backend_type {
        BackendType::Cpu => {
            if let Some(ref device_name) = info.device_name {
                format!(
                    "CPU ({})",
                    device_name
                        .split_whitespace()
                        .take(4)
                        .collect::<Vec<_>>()
                        .join(" ")
                )
            } else {
                "CPU".to_string()
            }
        }
        BackendType::Cuda => {
            if let Some(ref device_name) = info.device_name {
                format!("CUDA ({device_name})")
            } else {
                "CUDA".to_string()
            }
        }
        BackendType::Metal => {
            if let Some(ref device_name) = info.device_name {
                format!("Metal ({device_name})")
            } else {
                "Metal".to_string()
            }
        }
        BackendType::Rocm => {
            if let Some(ref device_name) = info.device_name {
                format!("ROCm ({device_name})")
            } else {
                "ROCm".to_string()
            }
        }
        _ => backend_type.to_string(),
    }
}

/// Get a human-readable device string for the current best available backend.
fn get_device_string() -> String {
    let backend_type = select_best_backend();
    get_device_string_for_backend(backend_type)
}

// ============================================================================
// AdamW Optimizer
// ============================================================================

/// `AdamW` optimizer with decoupled weight decay.
///
/// Implements the `AdamW` algorithm for parameter optimization with:
/// - Adaptive learning rates per parameter
/// - Momentum (first moment) tracking
/// - Variance (second moment) tracking
/// - Decoupled weight decay for better generalization
#[derive(Debug)]
struct AdamW {
    /// First moment estimates (momentum)
    m: Vec<Vec<f32>>,
    /// Second moment estimates (variance)
    v: Vec<Vec<f32>>,
    /// Timestep counter for bias correction
    t: u64,
    /// Learning rate
    lr: f32,
    /// First moment decay rate (β₁)
    beta1: f32,
    /// Second moment decay rate (β₂)
    beta2: f32,
    /// Numerical stability constant (ε)
    eps: f32,
    /// Weight decay coefficient (λ)
    weight_decay: f32,
}

impl AdamW {
    /// Create a new `AdamW` optimizer.
    ///
    /// # Arguments
    /// * `param_sizes` - Size of each parameter vector
    /// * `lr` - Learning rate (typical: 1e-3)
    /// * `beta1` - First moment decay (typical: 0.9)
    /// * `beta2` - Second moment decay (typical: 0.999)
    /// * `weight_decay` - Weight decay coefficient (typical: 0.01)
    fn new(param_sizes: &[usize], lr: f32, beta1: f32, beta2: f32, weight_decay: f32) -> Self {
        let m = param_sizes.iter().map(|&size| vec![0.0f32; size]).collect();
        let v = param_sizes.iter().map(|&size| vec![0.0f32; size]).collect();

        Self {
            m,
            v,
            t: 0,
            lr,
            beta1,
            beta2,
            eps: 1e-8,
            weight_decay,
        }
    }

    /// Create `AdamW` with default hyperparameters.
    ///
    /// Uses lr=1e-3, β₁=0.9, β₂=0.999, `weight_decay=0.01`
    #[allow(dead_code)]
    fn with_defaults(param_sizes: &[usize]) -> Self {
        Self::new(param_sizes, 1e-3, 0.9, 0.999, 0.01)
    }

    /// Perform one optimization step.
    ///
    /// Updates parameters in-place using `AdamW` update rule:
    /// ```text
    /// m = β₁ * m + (1 - β₁) * g           // Update momentum
    /// v = β₂ * v + (1 - β₂) * g²          // Update variance
    /// m̂ = m / (1 - β₁^t)                  // Bias-corrected momentum
    /// v̂ = v / (1 - β₂^t)                  // Bias-corrected variance
    /// θ = θ - lr * (m̂ / (√v̂ + ε) + λ * θ) // Decoupled weight decay update
    /// ```
    fn step(&mut self, params: &mut [Vec<f32>], grads: &[Vec<f32>]) {
        self.t += 1;

        // Bias correction factors
        let bias_correction1 = 1.0 - self.beta1.powi(self.t as i32);
        let bias_correction2 = 1.0 - self.beta2.powi(self.t as i32);

        for (i, (param, grad)) in params.iter_mut().zip(grads.iter()).enumerate() {
            for (j, (p, &g)) in param.iter_mut().zip(grad.iter()).enumerate() {
                // Update first moment: m = β₁ * m + (1 - β₁) * g
                self.m[i][j] = self.beta1 * self.m[i][j] + (1.0 - self.beta1) * g;

                // Update second moment: v = β₂ * v + (1 - β₂) * g²
                self.v[i][j] = self.beta2 * self.v[i][j] + (1.0 - self.beta2) * g * g;

                // Bias-corrected estimates
                let m_hat = self.m[i][j] / bias_correction1;
                let v_hat = self.v[i][j] / bias_correction2;

                // AdamW update with decoupled weight decay
                let adam_update = m_hat / (v_hat.sqrt() + self.eps);
                *p -= self.lr * (adam_update + self.weight_decay * *p);
            }
        }
    }

    /// Get current learning rate
    #[allow(dead_code)]
    fn learning_rate(&self) -> f32 {
        self.lr
    }

    /// Set learning rate
    #[allow(dead_code)]
    fn set_learning_rate(&mut self, lr: f32) {
        self.lr = lr;
    }
}

/// Extract all trainable parameters from `FNO2dWeights` as flat vectors.
fn extract_params(weights: &FNO2dWeights) -> Vec<Vec<f32>> {
    let mut params = Vec::new();

    // Lifting MLP parameters
    params.push(weights.lifting_fc1_weight.clone());
    params.push(weights.lifting_fc1_bias.clone());
    params.push(weights.lifting_fc2_weight.clone());
    params.push(weights.lifting_fc2_bias.clone());

    // Spectral conv parameters (for each layer)
    for sw in &weights.spectral_weights {
        params.push(sw.weights1_real.clone());
        params.push(sw.weights1_imag.clone());
        params.push(sw.weights2_real.clone());
        params.push(sw.weights2_imag.clone());
    }

    // 1x1 conv parameters
    for (weight, bias) in &weights.conv_weights {
        params.push(weight.clone());
        params.push(bias.clone());
    }

    // Projection parameters
    for (weight, bias) in &weights.projection_weights {
        params.push(weight.clone());
        params.push(bias.clone());
    }

    params
}

/// Update `FNO2dWeights` with new parameter values.
fn update_weights(weights: &mut FNO2dWeights, params: &[Vec<f32>]) {
    let mut idx = 0;

    // Lifting MLP parameters
    weights.lifting_fc1_weight = params[idx].clone();
    idx += 1;
    weights.lifting_fc1_bias = params[idx].clone();
    idx += 1;
    weights.lifting_fc2_weight = params[idx].clone();
    idx += 1;
    weights.lifting_fc2_bias = params[idx].clone();
    idx += 1;

    // Spectral conv parameters
    for sw in &mut weights.spectral_weights {
        sw.weights1_real = params[idx].clone();
        idx += 1;
        sw.weights1_imag = params[idx].clone();
        idx += 1;
        sw.weights2_real = params[idx].clone();
        idx += 1;
        sw.weights2_imag = params[idx].clone();
        idx += 1;
    }

    // 1x1 conv parameters
    for (weight, bias) in &mut weights.conv_weights {
        *weight = params[idx].clone();
        idx += 1;
        *bias = params[idx].clone();
        idx += 1;
    }

    // Projection parameters
    for (weight, bias) in &mut weights.projection_weights {
        *weight = params[idx].clone();
        idx += 1;
        *bias = params[idx].clone();
        idx += 1;
    }
}

// ============================================================================
// Early Stopping
// ============================================================================

/// Early stopping monitor to prevent overfitting.
///
/// Tracks validation loss and triggers early stopping when no improvement
/// is seen for a specified number of epochs (patience).
#[derive(Debug)]
struct EarlyStopping {
    /// Number of epochs to wait before stopping
    patience: usize,
    /// Minimum improvement required to reset counter
    min_delta: f32,
    /// Best validation loss seen so far
    best_loss: f32,
    /// Epochs since last improvement
    epochs_without_improvement: usize,
}

impl EarlyStopping {
    /// Create a new early stopping monitor.
    ///
    /// # Arguments
    /// * `patience` - Number of epochs to wait for improvement
    /// * `min_delta` - Minimum improvement required (default: 1e-4)
    fn new(patience: usize, min_delta: f32) -> Self {
        Self {
            patience,
            min_delta,
            best_loss: f32::INFINITY,
            epochs_without_improvement: 0,
        }
    }

    /// Check if training should stop.
    ///
    /// Returns `true` if no improvement has been seen for `patience` epochs.
    fn check(&mut self, val_loss: f32) -> bool {
        if val_loss < self.best_loss - self.min_delta {
            // Improvement found
            self.best_loss = val_loss;
            self.epochs_without_improvement = 0;
            false
        } else {
            // No improvement
            self.epochs_without_improvement += 1;
            self.epochs_without_improvement >= self.patience
        }
    }

    /// Get the best loss seen so far
    #[allow(dead_code)]
    fn best_loss(&self) -> f32 {
        self.best_loss
    }

    /// Get epochs without improvement
    fn epochs_without_improvement(&self) -> usize {
        self.epochs_without_improvement
    }
}

// ============================================================================
// Learning Rate Scheduler
// ============================================================================

/// Learning rate scheduling strategy.
#[derive(Debug, Clone)]
pub enum LRSchedulerType {
    /// Reduce LR when validation loss plateaus
    ReduceOnPlateau {
        /// Factor to reduce LR by (e.g., 0.1 = reduce to 10%)
        factor: f32,
        /// Epochs to wait before reducing
        patience: usize,
        /// Minimum learning rate
        min_lr: f32,
    },
    /// Cosine annealing from initial LR to min LR
    CosineAnnealing {
        /// Total epochs for one cycle
        t_max: usize,
        /// Minimum learning rate
        eta_min: f32,
    },
    /// Step decay: reduce LR every N epochs
    StepDecay {
        /// Reduce LR every `step_size` epochs
        step_size: usize,
        /// Factor to multiply LR by
        gamma: f32,
    },
}

/// Learning rate scheduler.
#[derive(Debug)]
pub struct LRScheduler {
    /// Scheduling strategy
    scheduler_type: LRSchedulerType,
    /// Initial learning rate
    initial_lr: f32,
    /// Current learning rate
    current_lr: f32,
    /// Best validation loss (for `ReduceOnPlateau`)
    best_val_loss: f32,
    /// Epochs since improvement (for `ReduceOnPlateau`)
    epochs_since_improvement: usize,
}

impl LRScheduler {
    /// Create a new LR scheduler.
    #[must_use]
    pub fn new(scheduler_type: LRSchedulerType, initial_lr: f32) -> Self {
        Self {
            scheduler_type,
            initial_lr,
            current_lr: initial_lr,
            best_val_loss: f32::INFINITY,
            epochs_since_improvement: 0,
        }
    }

    /// Update learning rate based on current epoch and validation loss.
    ///
    /// Returns the new learning rate.
    pub fn step(&mut self, epoch: usize, val_loss: Option<f32>) -> f32 {
        match &self.scheduler_type {
            LRSchedulerType::ReduceOnPlateau {
                factor,
                patience,
                min_lr,
            } => {
                if let Some(loss) = val_loss {
                    if loss < self.best_val_loss {
                        self.best_val_loss = loss;
                        self.epochs_since_improvement = 0;
                    } else {
                        self.epochs_since_improvement += 1;
                        if self.epochs_since_improvement >= *patience {
                            self.current_lr = (self.current_lr * factor).max(*min_lr);
                            self.epochs_since_improvement = 0;
                            tracing::info!("LR reduced to {:.2e}", self.current_lr);
                        }
                    }
                }
            }
            LRSchedulerType::CosineAnnealing { t_max, eta_min } => {
                // Cosine annealing: lr = eta_min + 0.5 * (initial_lr - eta_min) * (1 + cos(pi * t / T))
                let t = epoch % t_max;
                let cos_term = (std::f32::consts::PI * t as f32 / *t_max as f32).cos();
                self.current_lr = eta_min + 0.5 * (self.initial_lr - eta_min) * (1.0 + cos_term);
            }
            LRSchedulerType::StepDecay { step_size, gamma } => {
                // Reduce LR every step_size epochs
                let decay_count = epoch / step_size;
                self.current_lr = self.initial_lr * gamma.powi(decay_count as i32);
            }
        }
        self.current_lr
    }

    /// Get current learning rate
    #[must_use]
    pub fn learning_rate(&self) -> f32 {
        self.current_lr
    }
}

/// Convert a slice of f32 values to bytes (little-endian).
fn f32_to_bytes(values: &[f32]) -> Vec<u8> {
    values.iter().flat_map(|v| v.to_le_bytes()).collect()
}

/// Result type for training operations
pub type TrainingResult<T> = std::result::Result<T, TrainingError>;

/// Training-specific errors
#[derive(Debug, Clone)]
pub enum TrainingError {
    /// Training was cancelled
    Cancelled,
    /// Data generation failed
    DataGenerationFailed(String),
    /// Model initialization failed
    ModelInitFailed(String),
    /// Training step failed
    TrainingStepFailed(String),
    /// Weight save failed
    WeightSaveFailed(String),
    /// Already training
    AlreadyTraining,
    /// Not initialized
    NotInitialized,
}

impl std::fmt::Display for TrainingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Cancelled => write!(f, "Training was cancelled"),
            Self::DataGenerationFailed(msg) => write!(f, "Data generation failed: {msg}"),
            Self::ModelInitFailed(msg) => write!(f, "Model initialization failed: {msg}"),
            Self::TrainingStepFailed(msg) => write!(f, "Training step failed: {msg}"),
            Self::WeightSaveFailed(msg) => write!(f, "Failed to save weights: {msg}"),
            Self::AlreadyTraining => write!(f, "Training is already in progress"),
            Self::NotInitialized => write!(f, "Trainer not initialized"),
        }
    }
}

impl std::error::Error for TrainingError {}

/// Trainer for FNO models
///
/// Manages the full training lifecycle including data generation,
/// model training, and progress reporting.
///
/// This trainer is generic over the compute backend, allowing training
/// on CPU, CUDA, Metal, or `ROCm` devices.
pub struct FnoTrainer<B: Backend<FloatElem = f32>> {
    /// Training configuration
    config: TrainingConfig,
    /// PDE configuration
    pde_config: PDEConfig,
    /// Progress broadcast channel
    progress_tx: broadcast::Sender<TrainingProgress>,
    /// Flag indicating if training is active
    is_training: Arc<AtomicBool>,
    /// Flag to request cancellation
    cancel_requested: Arc<AtomicBool>,
    /// Current progress (for polling)
    current_progress: std::sync::RwLock<TrainingProgress>,
    /// Compute device
    device: B::Device,
}

impl<B: Backend<FloatElem = f32>> FnoTrainer<B>
where
    B::Device: Default,
{
    /// Create a new FNO trainer with the default device for this backend
    #[must_use]
    pub fn new(config: TrainingConfig, pde_config: PDEConfig) -> Self {
        Self::with_device(config, pde_config, B::Device::default())
    }
}

impl<B: Backend<FloatElem = f32>> FnoTrainer<B> {
    /// Get the device string for this backend type
    fn backend_device_string() -> String {
        // Use type_name to determine backend at runtime
        let type_name = std::any::type_name::<B>();

        if type_name.contains("CpuBackend") {
            get_device_string_for_backend(BackendType::Cpu)
        } else if type_name.contains("CudaBackend") {
            get_device_string_for_backend(BackendType::Cuda)
        } else if type_name.contains("MetalBackend") {
            get_device_string_for_backend(BackendType::Metal)
        } else if type_name.contains("RocmBackend") {
            get_device_string_for_backend(BackendType::Rocm)
        } else {
            // Fallback to global detection
            get_device_string()
        }
    }

    /// Create a new FNO trainer with a specific device
    #[must_use]
    pub fn with_device(config: TrainingConfig, pde_config: PDEConfig, device: B::Device) -> Self {
        let (progress_tx, _) = broadcast::channel(32);
        let n_batches = config.n_train_samples.div_ceil(config.batch_size);

        // Initialize progress with device info and learning rate
        let mut initial_progress = TrainingProgress::new(config.epochs, n_batches);
        initial_progress.device = Self::backend_device_string();
        initial_progress.current_lr = config.learning_rate;

        Self {
            config,
            pde_config,
            progress_tx,
            is_training: Arc::new(AtomicBool::new(false)),
            cancel_requested: Arc::new(AtomicBool::new(false)),
            current_progress: std::sync::RwLock::new(initial_progress),
            device,
        }
    }

    /// Check if training is currently active
    #[must_use]
    pub fn is_training(&self) -> bool {
        self.is_training.load(Ordering::SeqCst)
    }

    /// Subscribe to progress updates
    #[must_use]
    pub fn subscribe(&self) -> broadcast::Receiver<TrainingProgress> {
        self.progress_tx.subscribe()
    }

    /// Get current progress (for polling)
    #[must_use]
    pub fn get_progress(&self) -> TrainingProgress {
        self.current_progress.read().unwrap().clone()
    }

    /// Request training cancellation
    pub fn cancel(&self) {
        self.cancel_requested.store(true, Ordering::SeqCst);
    }

    /// Run training synchronously (blocking)
    ///
    /// Returns the path to saved weights on success.
    pub fn train(&self) -> TrainingResult<String> {
        if self
            .is_training
            .compare_exchange(false, true, Ordering::SeqCst, Ordering::SeqCst)
            .is_err()
        {
            return Err(TrainingError::AlreadyTraining);
        }

        // Reset cancellation flag
        self.cancel_requested.store(false, Ordering::SeqCst);

        let result = self.run_training_loop();

        // Mark training as complete
        self.is_training.store(false, Ordering::SeqCst);

        result
    }

    /// Run the main training loop
    fn run_training_loop(&self) -> TrainingResult<String> {
        let start_time = Instant::now();

        // Phase 1: Generate training data
        self.update_progress(|p| {
            p.status = TrainingStatus::GeneratingData {
                samples_generated: 0,
                total_samples: self.config.n_train_samples + self.config.n_val_samples,
            };
        });

        let data_config =
            DarcyDataConfig::default().with_resolution(self.pde_config.resolution as usize);

        let generator = DarcyDataGenerator::with_config(data_config);

        // Generate training samples
        let train_samples =
            self.generate_samples_with_progress(&generator, self.config.n_train_samples, 0)?;

        // Generate validation samples
        let val_samples = self.generate_samples_with_progress(
            &generator,
            self.config.n_val_samples,
            self.config.n_train_samples,
        )?;

        if self.cancel_requested.load(Ordering::SeqCst) {
            self.update_progress(|p| p.status = TrainingStatus::Cancelled);
            return Err(TrainingError::Cancelled);
        }

        // Phase 2: Initialize model and optimizer
        self.update_progress(|p| p.status = TrainingStatus::Training);

        let pde_name = self
            .pde_config
            .pde_type
            .name()
            .to_lowercase()
            .replace(' ', "_");

        // Create initial model (using generic backend B)
        let model = FNO2d::<B>::new_with_layers(
            1, // in_channels (permeability field)
            1, // out_channels (pressure field)
            self.pde_config.model_width as usize,
            self.pde_config.n_modes.0 as usize,
            self.pde_config.n_layers as usize,
            &self.device,
        )
        .map_err(|e| TrainingError::ModelInitFailed(e.to_string()))?;

        // Extract initial weights and create optimizer
        let mut weights = model.to_weights(&pde_name);
        let mut params = extract_params(&weights);
        let param_sizes: Vec<usize> = params.iter().map(std::vec::Vec::len).collect();

        // Initialize AdamW optimizer with user-configured learning rate
        let mut optimizer = AdamW::new(
            &param_sizes,
            self.config.learning_rate,
            0.9,   // beta1
            0.999, // beta2
            0.01,  // weight_decay
        );

        tracing::info!(
            "Initialized AdamW optimizer: lr={}, params={}, total_elements={}",
            optimizer.learning_rate(),
            params.len(),
            param_sizes.iter().sum::<usize>()
        );

        // Phase 3: Training loop with AdamW optimization
        let n_batches = train_samples.len().div_ceil(self.config.batch_size);
        let mut best_loss = f32::MAX;
        let mut loss_history = Vec::with_capacity(self.config.epochs);

        // Initialize early stopping (patience = 10 epochs, min_delta = 1e-4)
        let mut early_stopping = EarlyStopping::new(10, 1e-4);

        // Initialize LR scheduler (ReduceOnPlateau with factor=0.5, patience=5)
        let mut lr_scheduler = LRScheduler::new(
            LRSchedulerType::ReduceOnPlateau {
                factor: 0.5,
                patience: 5,
                min_lr: 1e-6,
            },
            self.config.learning_rate,
        );

        // Epsilon for numerical gradient computation
        const GRAD_EPS: f32 = 1e-4;

        for epoch in 0..self.config.epochs {
            if self.cancel_requested.load(Ordering::SeqCst) {
                self.update_progress(|p| p.status = TrainingStatus::Cancelled);
                return Err(TrainingError::Cancelled);
            }

            let epoch_start = Instant::now();
            let mut epoch_loss = 0.0f32;

            // Process batches
            for batch_idx in 0..n_batches {
                let batch_start = batch_idx * self.config.batch_size;
                let batch_end = (batch_start + self.config.batch_size).min(train_samples.len());
                let batch_samples = &train_samples[batch_start..batch_end];

                // Create batch tensors
                let batch = TrainingBatch::from_samples(batch_samples);

                // Rebuild model with current parameters
                update_weights(&mut weights, &params);
                let current_model = FNO2d::<B>::from_weights(&weights, &self.device)
                    .map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;

                // Compute current loss
                let batch_loss = self.compute_batch_loss(&current_model, &batch)?;
                epoch_loss += batch_loss;

                // Compute gradients using numerical differentiation
                // For efficiency, we use stochastic parameter sampling
                let grads = self.compute_numerical_gradients(
                    &mut weights,
                    &params,
                    &batch,
                    batch_loss,
                    GRAD_EPS,
                )?;

                // Apply AdamW optimizer step
                optimizer.step(&mut params, &grads);

                // Update progress
                let elapsed = start_time.elapsed().as_secs_f32();
                let samples_processed = epoch * train_samples.len() + batch_end;
                let total_samples = self.config.epochs * train_samples.len();
                let progress_frac = samples_processed as f32 / total_samples as f32;
                let eta = if progress_frac > 0.01 {
                    elapsed / progress_frac - elapsed
                } else {
                    0.0
                };

                // Capture current learning rate for progress reporting
                let current_lr = optimizer.lr;

                self.update_progress(|p| {
                    p.epoch = epoch;
                    p.batch = batch_idx;
                    p.total_batches = n_batches;
                    p.elapsed_seconds = elapsed;
                    p.eta_seconds = eta;
                    p.samples_per_sec = samples_processed as f32 / elapsed.max(0.01);
                    p.current_lr = current_lr;
                });
            }

            // Compute average epoch loss
            let avg_loss = epoch_loss / n_batches as f32;
            loss_history.push(avg_loss);

            if avg_loss < best_loss {
                best_loss = avg_loss;
            }

            // Rebuild model for validation
            update_weights(&mut weights, &params);
            let val_model = FNO2d::<B>::from_weights(&weights, &self.device)
                .map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;

            // Validation loss
            let val_loss = self.compute_validation_loss(&val_model, &val_samples)?;

            // Update LR scheduler based on validation loss
            let new_lr = lr_scheduler.step(epoch, Some(val_loss));
            optimizer.lr = new_lr;

            // Update progress with epoch results
            let current_lr = new_lr;
            let epochs_without_improvement = early_stopping.epochs_without_improvement();
            self.update_progress(|p| {
                p.loss = avg_loss;
                p.best_loss = best_loss;
                p.val_loss = Some(val_loss);
                p.loss_history = loss_history.clone();
                p.current_lr = current_lr;
            });

            // Log progress
            let epoch_time = epoch_start.elapsed().as_secs_f32();
            tracing::info!(
                "Epoch {}/{}: loss={:.6}, val_loss={:.6}, time={:.2}s, lr={:.2e}, patience={}/10",
                epoch + 1,
                self.config.epochs,
                avg_loss,
                val_loss,
                epoch_time,
                new_lr,
                epochs_without_improvement
            );

            // Check early stopping
            if early_stopping.check(val_loss) {
                tracing::info!(
                    "Early stopping triggered at epoch {} (no improvement for {} epochs)",
                    epoch + 1,
                    early_stopping.epochs_without_improvement()
                );
                break;
            }
        }

        // Phase 4: Save final weights
        update_weights(&mut weights, &params);
        let final_model = FNO2d::<B>::from_weights(&weights, &self.device)
            .map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
        let weights_path = self.save_weights(&final_model)?;

        // Mark complete
        self.update_progress(|p| {
            p.status = TrainingStatus::Complete;
            p.epoch = self.config.epochs;
        });

        Ok(weights_path)
    }

    /// Generate samples with progress tracking
    fn generate_samples_with_progress(
        &self,
        generator: &DarcyDataGenerator,
        n_samples: usize,
        offset: usize,
    ) -> TrainingResult<Vec<DarcySample>> {
        let total = self.config.n_train_samples + self.config.n_val_samples;

        let samples = generator.generate_samples_with_progress(n_samples, |current, _| {
            if current % 10 == 0 || current == n_samples {
                self.update_progress(|p| {
                    p.status = TrainingStatus::GeneratingData {
                        samples_generated: offset + current,
                        total_samples: total,
                    };
                });
            }
        });

        Ok(samples)
    }

    /// Compute loss for a single batch
    fn compute_batch_loss(&self, model: &FNO2d<B>, batch: &TrainingBatch) -> TrainingResult<f32> {
        let batch_size = batch.batch_size;
        let resolution = batch.resolution;

        // Create input tensor [B, 1, H, W]
        let input: GenericTensor<B, 4> = GenericTensor::from_slice(
            &batch.inputs,
            [batch_size, 1, resolution, resolution],
            &self.device,
        );

        // Forward pass
        let output = model.forward_4d(&input);
        let output_data = output.to_vec();

        // Compute MSE loss
        let mut loss = 0.0f32;
        for (pred, target) in output_data.iter().zip(batch.outputs.iter()) {
            let diff = pred - target;
            loss += diff * diff;
        }

        Ok(loss / batch.total_elements() as f32)
    }

    /// Compute validation loss
    fn compute_validation_loss(
        &self,
        model: &FNO2d<B>,
        samples: &[DarcySample],
    ) -> TrainingResult<f32> {
        if samples.is_empty() {
            return Ok(0.0);
        }

        let batch = TrainingBatch::from_samples(samples);
        self.compute_batch_loss(model, &batch)
    }

    /// Compute numerical gradients using stochastic parameter perturbation.
    ///
    /// For efficiency, we sample a subset of parameters to compute gradients
    /// rather than computing gradients for all parameters every step.
    /// This provides an unbiased gradient estimate while being computationally feasible.
    ///
    /// # Arguments
    /// * `weights` - Model weights structure (will be temporarily modified)
    /// * `params` - Current parameter values
    /// * `batch` - Training batch
    /// * `base_loss` - Loss computed at current parameter values
    /// * `eps` - Perturbation epsilon for finite differences
    #[allow(clippy::too_many_arguments)]
    fn compute_numerical_gradients(
        &self,
        weights: &mut FNO2dWeights,
        params: &[Vec<f32>],
        batch: &TrainingBatch,
        base_loss: f32,
        eps: f32,
    ) -> TrainingResult<Vec<Vec<f32>>> {
        // Initialize gradients with zeros
        let mut grads: Vec<Vec<f32>> = params.iter().map(|p| vec![0.0f32; p.len()]).collect();

        // For efficiency, we use coordinate descent with stochastic sampling:
        // Sample a subset of parameters per batch for gradient computation.
        // This provides unbiased gradient estimates over time.

        // Maximum number of parameter elements to sample per batch
        const MAX_SAMPLES_PER_BATCH: usize = 256;

        // Count total parameters
        let total_params: usize = params.iter().map(std::vec::Vec::len).sum();

        // Use a simple hash of the loss to seed randomness (deterministic within batch)
        let seed = (base_loss.to_bits() ^ batch.batch_size as u32) as usize;

        // Compute how many samples we can afford
        let n_samples = MAX_SAMPLES_PER_BATCH.min(total_params);

        // Sample parameter indices
        let mut perturbed_params = params.to_vec();

        for sample_idx in 0..n_samples {
            // Deterministic "random" parameter selection using simple hashing
            let flat_idx = (seed.wrapping_mul(31).wrapping_add(sample_idx * 17)) % total_params;

            // Convert flat index to (param_group, param_idx)
            let mut remaining = flat_idx;
            let mut param_group = 0;
            while param_group < params.len() && remaining >= params[param_group].len() {
                remaining -= params[param_group].len();
                param_group += 1;
            }

            if param_group >= params.len() {
                continue;
            }

            let param_idx = remaining;

            // Compute gradient using central differences: (f(x+eps) - f(x-eps)) / (2*eps)
            // This is more accurate than forward differences

            // Perturb parameter positively
            let original_value = perturbed_params[param_group][param_idx];
            perturbed_params[param_group][param_idx] = original_value + eps;
            update_weights(weights, &perturbed_params);

            let perturbed_model = FNO2d::<B>::from_weights(weights, &self.device)
                .map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
            let loss_plus = self.compute_batch_loss(&perturbed_model, batch)?;

            // Perturb parameter negatively
            perturbed_params[param_group][param_idx] = original_value - eps;
            update_weights(weights, &perturbed_params);

            let perturbed_model = FNO2d::<B>::from_weights(weights, &self.device)
                .map_err(|e| TrainingError::TrainingStepFailed(e.to_string()))?;
            let loss_minus = self.compute_batch_loss(&perturbed_model, batch)?;

            // Restore original value
            perturbed_params[param_group][param_idx] = original_value;

            // Compute gradient using central differences
            let grad = (loss_plus - loss_minus) / (2.0 * eps);

            // Scale gradient by ratio of total to sampled parameters for unbiased estimate
            let scale = total_params as f32 / n_samples as f32;
            grads[param_group][param_idx] = grad * scale;
        }

        // Restore weights to original state
        update_weights(weights, params);

        Ok(grads)
    }

    /// Save trained weights to `SafeTensors` format
    fn save_weights(&self, model: &FNO2d<B>) -> TrainingResult<String> {
        use rtx_hub::safetensors::{SafeTensorsBuilder, SafeTensorsDType};

        // Create weights directory
        let weights_dir = dirs::data_local_dir()
            .unwrap_or_else(std::env::temp_dir)
            .join("rustytorch/neural-operator");
        std::fs::create_dir_all(&weights_dir)
            .map_err(|e| TrainingError::WeightSaveFailed(e.to_string()))?;

        let pde_name = self
            .pde_config
            .pde_type
            .name()
            .to_lowercase()
            .replace(' ', "_");
        let weights_path = weights_dir.join(format!(
            "fno_{}_{}.safetensors",
            pde_name, self.pde_config.resolution
        ));

        // Extract weights from model
        let weights = model.to_weights(&pde_name);

        // Build SafeTensors file
        let builder = SafeTensorsBuilder::new()
            // Metadata
            .with_metadata("model_type", "FNO2d")
            .with_metadata("pde_type", &weights.config.pde_type)
            .with_metadata("resolution", self.pde_config.resolution.to_string())
            .with_metadata("in_channels", weights.config.in_channels.to_string())
            .with_metadata("out_channels", weights.config.out_channels.to_string())
            .with_metadata("width", weights.config.width.to_string())
            .with_metadata("n_layers", weights.config.n_layers.to_string())
            .with_metadata("n_modes_h", weights.config.n_modes.0.to_string())
            .with_metadata("n_modes_w", weights.config.n_modes.1.to_string())
            // Lifting MLP weights
            .add_tensor(
                "lifting.fcs.0.weight",
                SafeTensorsDType::F32,
                vec![weights.config.width * 2, weights.config.in_channels + 2],
                f32_to_bytes(&weights.lifting_fc1_weight),
            )
            .add_tensor(
                "lifting.fcs.0.bias",
                SafeTensorsDType::F32,
                vec![weights.config.width * 2],
                f32_to_bytes(&weights.lifting_fc1_bias),
            )
            .add_tensor(
                "lifting.fcs.1.weight",
                SafeTensorsDType::F32,
                vec![weights.config.width, weights.config.width * 2],
                f32_to_bytes(&weights.lifting_fc2_weight),
            )
            .add_tensor(
                "lifting.fcs.1.bias",
                SafeTensorsDType::F32,
                vec![weights.config.width],
                f32_to_bytes(&weights.lifting_fc2_bias),
            );

        // Add spectral conv weights for each layer
        let mut builder = builder;
        for (i, sw) in weights.spectral_weights.iter().enumerate() {
            let shape = vec![
                weights.config.width,
                weights.config.width,
                weights.config.n_modes.0,
                weights.config.n_modes.1,
            ];
            builder = builder
                .add_tensor(
                    format!("spectral_conv.{i}.weights1_real"),
                    SafeTensorsDType::F32,
                    shape.clone(),
                    f32_to_bytes(&sw.weights1_real),
                )
                .add_tensor(
                    format!("spectral_conv.{i}.weights1_imag"),
                    SafeTensorsDType::F32,
                    shape.clone(),
                    f32_to_bytes(&sw.weights1_imag),
                )
                .add_tensor(
                    format!("spectral_conv.{i}.weights2_real"),
                    SafeTensorsDType::F32,
                    shape.clone(),
                    f32_to_bytes(&sw.weights2_real),
                )
                .add_tensor(
                    format!("spectral_conv.{i}.weights2_imag"),
                    SafeTensorsDType::F32,
                    shape,
                    f32_to_bytes(&sw.weights2_imag),
                );
        }

        // Add 1x1 conv weights for each layer
        for (i, (weight, bias)) in weights.conv_weights.iter().enumerate() {
            builder = builder
                .add_tensor(
                    format!("conv.{i}.weight"),
                    SafeTensorsDType::F32,
                    vec![weights.config.width, weights.config.width],
                    f32_to_bytes(weight),
                )
                .add_tensor(
                    format!("conv.{i}.bias"),
                    SafeTensorsDType::F32,
                    vec![weights.config.width],
                    f32_to_bytes(bias),
                );
        }

        // Add projection layer weights
        const PROJECTION_HIDDEN: usize = 128;
        let (proj0_weight, proj0_bias) = &weights.projection_weights[0];
        let (proj1_weight, proj1_bias) = &weights.projection_weights[1];

        builder = builder
            .add_tensor(
                "projection.0.weight",
                SafeTensorsDType::F32,
                vec![PROJECTION_HIDDEN, weights.config.width],
                f32_to_bytes(proj0_weight),
            )
            .add_tensor(
                "projection.0.bias",
                SafeTensorsDType::F32,
                vec![PROJECTION_HIDDEN],
                f32_to_bytes(proj0_bias),
            )
            .add_tensor(
                "projection.1.weight",
                SafeTensorsDType::F32,
                vec![weights.config.out_channels, PROJECTION_HIDDEN],
                f32_to_bytes(proj1_weight),
            )
            .add_tensor(
                "projection.1.bias",
                SafeTensorsDType::F32,
                vec![weights.config.out_channels],
                f32_to_bytes(proj1_bias),
            );

        // Build and write file
        let file_bytes = builder
            .build()
            .map_err(|e| TrainingError::WeightSaveFailed(e.to_string()))?;

        std::fs::write(&weights_path, file_bytes)
            .map_err(|e| TrainingError::WeightSaveFailed(e.to_string()))?;

        tracing::info!("Weights saved to {:?}", weights_path);

        Ok(weights_path.to_string_lossy().to_string())
    }

    /// Update progress and broadcast
    fn update_progress<F>(&self, updater: F)
    where
        F: FnOnce(&mut TrainingProgress),
    {
        let mut progress = self.current_progress.write().unwrap();
        updater(&mut progress);

        // Broadcast update (ignore send errors - no receivers)
        let _ = self.progress_tx.send(progress.clone());
    }
}

// ============================================================================
// Dynamic Backend Dispatch
// ============================================================================

/// Trait object for dynamic dispatch of training operations.
///
/// This enables runtime selection of compute backend without requiring
/// the caller to know the specific backend type.
pub trait DynamicTrainer: Send + Sync {
    /// Start training and return the weights path on success
    fn train(&self) -> TrainingResult<String>;
    /// Get current training progress
    fn get_progress(&self) -> TrainingProgress;
    /// Cancel training
    fn cancel(&self);
    /// Check if training is active
    fn is_training(&self) -> bool;
}

impl<B: Backend<FloatElem = f32> + Send + Sync + 'static> DynamicTrainer for FnoTrainer<B>
where
    B::Device: Send + Sync,
{
    fn train(&self) -> TrainingResult<String> {
        FnoTrainer::train(self)
    }

    fn get_progress(&self) -> TrainingProgress {
        FnoTrainer::get_progress(self)
    }

    fn cancel(&self) {
        FnoTrainer::cancel(self);
    }

    fn is_training(&self) -> bool {
        FnoTrainer::is_training(self)
    }
}

/// Training session handle for async operations.
///
/// This struct uses runtime backend selection to automatically choose
/// the best available compute device (CUDA > Metal > `ROCm` > CPU).
pub struct TrainingSession {
    trainer: Arc<dyn DynamicTrainer>,
}

impl TrainingSession {
    /// Create a new training session with automatic backend selection.
    ///
    /// The backend is selected based on availability:
    /// 1. CUDA (if nvidia-smi available and feature enabled)
    /// 2. Metal (on macOS, if feature enabled)
    /// 3. `ROCm` (if rocminfo available and feature enabled)
    /// 4. CPU (always available)
    pub fn new(config: TrainingConfig, pde_config: PDEConfig) -> Self {
        let backend_type = select_best_backend();
        tracing::info!("Training session using backend: {}", backend_type);

        let trainer: Arc<dyn DynamicTrainer> = match backend_type {
            #[cfg(feature = "cuda")]
            BackendType::Cuda => {
                tracing::info!("Initializing CUDA backend for training");
                Arc::new(FnoTrainer::<CudaBackend>::new(config, pde_config))
            }
            #[cfg(feature = "metal")]
            BackendType::Metal => {
                tracing::info!("Initializing Metal backend for training");
                Arc::new(FnoTrainer::<MetalBackend>::new(config, pde_config))
            }
            #[cfg(feature = "rocm")]
            BackendType::Rocm => {
                tracing::info!("Initializing ROCm backend for training");
                Arc::new(FnoTrainer::<RocmBackend>::new(config, pde_config))
            }
            _ => {
                tracing::info!("Initializing CPU backend for training");
                Arc::new(FnoTrainer::<CpuBackend>::new(config, pde_config))
            }
        };

        Self { trainer }
    }

    /// Create a training session with a specific backend.
    ///
    /// This is useful for testing or when you want to force a specific backend.
    #[must_use]
    pub fn with_backend<B: Backend<FloatElem = f32> + Send + Sync + 'static>(
        config: TrainingConfig,
        pde_config: PDEConfig,
    ) -> Self
    where
        B::Device: Default + Send + Sync,
    {
        Self {
            trainer: Arc::new(FnoTrainer::<B>::new(config, pde_config)),
        }
    }

    /// Start training in background
    #[must_use]
    pub fn start(&self) -> tokio::task::JoinHandle<TrainingResult<String>> {
        let trainer = Arc::clone(&self.trainer);
        tokio::task::spawn_blocking(move || trainer.train())
    }

    /// Get current progress
    #[must_use]
    pub fn get_progress(&self) -> TrainingProgress {
        self.trainer.get_progress()
    }

    /// Cancel training
    pub fn cancel(&self) {
        self.trainer.cancel();
    }

    /// Check if training is active
    #[must_use]
    pub fn is_training(&self) -> bool {
        self.trainer.is_training()
    }
}

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

    #[test]
    fn test_trainer_creation() {
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(32);
        // Explicitly use CpuBackend for tests
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);

        assert!(!trainer.is_training());
    }

    #[test]
    fn test_progress_update() {
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(32);
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);

        let initial_progress = trainer.get_progress();
        assert_eq!(initial_progress.epoch, 0);
        assert_eq!(initial_progress.status, TrainingStatus::NotStarted);

        trainer.update_progress(|p| {
            p.epoch = 5;
            p.loss = 0.01;
        });

        let updated_progress = trainer.get_progress();
        assert_eq!(updated_progress.epoch, 5);
        assert!((updated_progress.loss - 0.01).abs() < 1e-6);
    }

    #[test]
    fn test_cancel_flag() {
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(32);
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);

        assert!(!trainer.cancel_requested.load(Ordering::SeqCst));
        trainer.cancel();
        assert!(trainer.cancel_requested.load(Ordering::SeqCst));
    }

    // ============================================================================
    // TDD Tests for AdamW Optimizer (RED-GREEN-REFACTOR)
    // ============================================================================

    #[test]
    fn test_adamw_optimizer_creation() {
        // RED: Test that AdamW optimizer can be created with config
        let param_sizes = vec![4, 2];
        let learning_rate = 0.001;
        let beta1 = 0.9;
        let beta2 = 0.999;
        let weight_decay = 0.01;

        let optimizer = AdamW::new(&param_sizes, learning_rate, beta1, beta2, weight_decay);

        // GREEN: Verify optimizer was created successfully
        assert_eq!(optimizer.lr, learning_rate);
        assert_eq!(optimizer.beta1, beta1);
        assert_eq!(optimizer.beta2, beta2);
        assert_eq!(optimizer.weight_decay, weight_decay);
        assert_eq!(optimizer.eps, 1e-8);
        assert_eq!(optimizer.t, 0, "Initial timestep should be 0");

        // Verify momentum and variance are initialized to zeros
        assert_eq!(optimizer.m.len(), param_sizes.len());
        assert_eq!(optimizer.v.len(), param_sizes.len());
        for (i, &size) in param_sizes.iter().enumerate() {
            assert_eq!(optimizer.m[i].len(), size);
            assert_eq!(optimizer.v[i].len(), size);
            assert!(optimizer.m[i].iter().all(|&x| x == 0.0));
            assert!(optimizer.v[i].iter().all(|&x| x == 0.0));
        }
    }

    #[test]
    fn test_adamw_optimizer_step_updates_weights() {
        // RED: Test that optimizer step updates weights
        let param_sizes = vec![4, 2];
        let mut optimizer = AdamW::new(&param_sizes, 0.1, 0.9, 0.999, 0.01);

        let mut params = vec![vec![1.0, 2.0, 3.0, 4.0], vec![5.0, 6.0]];

        let grads = vec![vec![0.1, 0.2, 0.3, 0.4], vec![0.5, 0.6]];

        let original_params = params.clone();

        // GREEN: Perform optimization step
        optimizer.step(&mut params, &grads);

        // All parameters should have changed
        for (i, (new_param, old_param)) in params.iter().zip(original_params.iter()).enumerate() {
            for (j, (new_val, old_val)) in new_param.iter().zip(old_param.iter()).enumerate() {
                assert!(
                    (new_val - old_val).abs() > 1e-8,
                    "Parameter [{i}][{j}] should have changed: {old_val} -> {new_val}"
                );
            }
        }

        // All parameters should have decreased (positive gradient with descent)
        for (i, (new_param, old_param)) in params.iter().zip(original_params.iter()).enumerate() {
            for (j, (new_val, old_val)) in new_param.iter().zip(old_param.iter()).enumerate() {
                assert!(
                    *new_val < *old_val,
                    "Parameter [{i}][{j}] should decrease: {old_val} -> {new_val}"
                );
            }
        }
    }

    #[test]
    fn test_adamw_maintains_momentum_state() {
        // RED: Test that momentum state is maintained across batches
        let param_sizes = vec![4];
        let mut optimizer = AdamW::new(&param_sizes, 0.01, 0.9, 0.999, 0.01);

        let mut params = vec![vec![1.0, 1.0, 1.0, 1.0]];
        let grads = vec![vec![0.1, 0.1, 0.1, 0.1]];

        // GREEN: First step
        optimizer.step(&mut params, &grads);

        // After first step, momentum should be non-zero
        assert!(
            optimizer.m[0].iter().any(|&x| x.abs() > 1e-8),
            "Momentum should be non-zero after first step"
        );

        // Expected momentum after first step: m = β₁ * 0 + (1 - β₁) * 0.1 = 0.1 * 0.1 = 0.01
        let expected_m_first = (1.0 - 0.9) * 0.1;
        assert!(
            (optimizer.m[0][0] - expected_m_first).abs() < 1e-6,
            "Momentum after first step should be {}, got {}",
            expected_m_first,
            optimizer.m[0][0]
        );

        // Second step with same gradient
        let m_before_second = optimizer.m[0][0];
        optimizer.step(&mut params, &grads);

        // Momentum should have accumulated: m = β₁ * m_prev + (1 - β₁) * g
        let expected_m_second = 0.9 * m_before_second + (1.0 - 0.9) * 0.1;
        assert!(
            (optimizer.m[0][0] - expected_m_second).abs() < 1e-6,
            "Momentum after second step should be {}, got {}",
            expected_m_second,
            optimizer.m[0][0]
        );
    }

    #[test]
    fn test_adamw_maintains_variance_state() {
        // RED: Test that variance state is maintained across batches
        let param_sizes = vec![4];
        let mut optimizer = AdamW::new(&param_sizes, 0.01, 0.9, 0.999, 0.01);

        let mut params = vec![vec![1.0, 1.0, 1.0, 1.0]];
        let grads = vec![vec![0.1, 0.1, 0.1, 0.1]];

        // GREEN: First step
        optimizer.step(&mut params, &grads);

        // After first step, variance should be non-zero
        assert!(
            optimizer.v[0].iter().any(|&x| x.abs() > 1e-8),
            "Variance should be non-zero after first step"
        );

        // Expected variance after first step: v = β₂ * 0 + (1 - β₂) * g² = 0.001 * 0.01 = 0.00001
        let expected_v_first = (1.0 - 0.999) * 0.1 * 0.1;
        assert!(
            (optimizer.v[0][0] - expected_v_first).abs() < 1e-8,
            "Variance after first step should be {}, got {}",
            expected_v_first,
            optimizer.v[0][0]
        );

        // Second step
        let v_before_second = optimizer.v[0][0];
        optimizer.step(&mut params, &grads);

        // Variance should have accumulated
        let expected_v_second = 0.999 * v_before_second + (1.0 - 0.999) * 0.1 * 0.1;
        assert!(
            (optimizer.v[0][0] - expected_v_second).abs() < 1e-8,
            "Variance after second step should be {}, got {}",
            expected_v_second,
            optimizer.v[0][0]
        );
    }

    #[test]
    fn test_adamw_timestep_increments() {
        // RED: Test that timestep counter increments with each step
        let param_sizes = vec![2];
        let mut optimizer = AdamW::new(&param_sizes, 0.01, 0.9, 0.999, 0.01);
        let mut params = vec![vec![1.0, 2.0]];
        let grads = vec![vec![0.1, 0.1]];

        // GREEN: Verify timestep increments
        assert_eq!(optimizer.t, 0, "Initial timestep should be 0");

        optimizer.step(&mut params, &grads);
        assert_eq!(optimizer.t, 1, "Timestep should be 1 after first step");

        optimizer.step(&mut params, &grads);
        assert_eq!(optimizer.t, 2, "Timestep should be 2 after second step");

        optimizer.step(&mut params, &grads);
        assert_eq!(optimizer.t, 3, "Timestep should be 3 after third step");
    }

    #[test]
    fn test_adamw_loss_decreases_faster_than_sgd() {
        // RED: Test that loss decreases faster with AdamW than with manual SGD
        // This demonstrates the adaptive learning rate benefit

        // Setup: Two identical training scenarios
        let param_sizes = vec![10]; // Single weight vector
        let mut params_adamw = vec![vec![1.0; 10]];
        let mut params_sgd = vec![vec![1.0; 10]];

        let learning_rate = 0.01;
        let mut optimizer_adamw = AdamW::new(&param_sizes, learning_rate, 0.9, 0.999, 0.0);

        // Simulate 20 steps with constant gradient
        let grads = vec![vec![0.1; 10]];

        // GREEN: Run optimization
        for _ in 0..20 {
            // AdamW step
            optimizer_adamw.step(&mut params_adamw, &grads);

            // Manual SGD step: param -= lr * grad
            for (param_vec, grad_vec) in params_sgd.iter_mut().zip(grads.iter()) {
                for (p, &g) in param_vec.iter_mut().zip(grad_vec.iter()) {
                    *p -= learning_rate * g;
                }
            }
        }

        // Compute "loss" as squared distance from origin (simulating convergence to zero)
        let loss_adamw: f32 = params_adamw[0].iter().map(|&x| x * x).sum();
        let loss_sgd: f32 = params_sgd[0].iter().map(|&x| x * x).sum();

        // AdamW should converge faster due to adaptive learning and momentum
        assert!(
            loss_adamw < loss_sgd,
            "AdamW loss ({}) should be lower than SGD loss ({}) after 20 steps",
            loss_adamw,
            loss_sgd
        );
    }

    #[test]
    fn test_adamw_bias_correction() {
        // RED: Test that bias correction is applied properly in early steps
        let param_sizes = vec![2];
        let mut optimizer = AdamW::new(&param_sizes, 1.0, 0.9, 0.999, 0.0);
        let mut params = vec![vec![1.0, 1.0]];
        let grads = vec![vec![1.0, 1.0]];

        let params_before = params.clone();

        // GREEN: First step
        optimizer.step(&mut params, &grads);

        // Without bias correction, the update would be very small
        // With bias correction: m_hat = m / (1 - β₁^t) = 0.1 / 0.1 = 1.0
        // The parameter change should be significant due to bias correction
        let change = (params[0][0] - params_before[0][0]).abs();
        assert!(
            change > 0.5,
            "First step should have large update due to bias correction, change = {}",
            change
        );
    }

    #[test]
    fn test_adamw_weight_decay_applied() {
        // RED: Test that weight decay is applied correctly
        let param_sizes = vec![4];
        let weight_decay = 0.1; // Significant weight decay for testing
        let mut optimizer = AdamW::new(&param_sizes, 0.001, 0.9, 0.999, weight_decay);

        let mut params = vec![vec![1.0, 2.0, 3.0, 4.0]];
        let grads = vec![vec![0.0, 0.0, 0.0, 0.0]]; // Zero gradient

        let params_before = params.clone();

        // GREEN: With zero gradient and weight decay, parameters should shrink toward zero
        optimizer.step(&mut params, &grads);

        for (i, (&new_val, &old_val)) in params[0].iter().zip(params_before[0].iter()).enumerate() {
            assert!(
                new_val < old_val,
                "Parameter [{}] should decrease due to weight decay: {} -> {}",
                i,
                old_val,
                new_val
            );
        }
    }

    #[test]
    fn test_adamw_with_default_hyperparameters() {
        // RED: Test that default hyperparameters produce reasonable behavior
        let param_sizes = vec![4, 2];
        let optimizer = AdamW::with_defaults(&param_sizes);

        // GREEN: Verify defaults
        assert_eq!(optimizer.lr, 1e-3);
        assert_eq!(optimizer.beta1, 0.9);
        assert_eq!(optimizer.beta2, 0.999);
        assert_eq!(optimizer.weight_decay, 0.01);
    }

    #[test]
    fn test_adamw_integration_with_training_config() {
        // RED: Test that AdamW integrates with training configuration
        let config = TrainingConfig {
            epochs: 2,
            batch_size: 4,
            learning_rate: 0.002, // Custom learning rate
            n_train_samples: 8,
            n_val_samples: 4,
        };

        let param_sizes = vec![10, 5];
        let optimizer = AdamW::new(&param_sizes, config.learning_rate, 0.9, 0.999, 0.01);

        // GREEN: Verify optimizer uses config learning rate
        assert_eq!(optimizer.lr, config.learning_rate);
    }

    #[test]
    fn test_extract_update_weights_roundtrip() {
        // Test that extract_params and update_weights are inverses
        let device = CpuDevice::default();
        let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
            .expect("Failed to create model");

        let mut weights = model.to_weights("test");
        let original_first_weight = weights.lifting_fc1_weight[0];

        let params = extract_params(&weights);

        // Verify extracted parameters match
        assert!(
            (params[0][0] - original_first_weight).abs() < 1e-6,
            "Extracted params should match original"
        );

        // Modify a parameter
        let mut modified_params = params.clone();
        modified_params[0][0] += 1.0;

        // Update weights and check the modification took effect
        update_weights(&mut weights, &modified_params);

        assert!(
            (weights.lifting_fc1_weight[0] - original_first_weight - 1.0).abs() < 1e-6,
            "Weight modification should persist"
        );
    }

    #[test]
    fn test_early_stopping() {
        let mut early_stop = EarlyStopping::new(3, 1e-4);

        // First check with good loss - should not stop
        assert!(!early_stop.check(1.0));
        assert_eq!(early_stop.epochs_without_improvement(), 0);

        // Improvement - counter stays at 0
        assert!(!early_stop.check(0.5));
        assert_eq!(early_stop.epochs_without_improvement(), 0);

        // No improvement - counter increases
        assert!(!early_stop.check(0.6));
        assert_eq!(early_stop.epochs_without_improvement(), 1);

        assert!(!early_stop.check(0.7));
        assert_eq!(early_stop.epochs_without_improvement(), 2);

        // Third no improvement - triggers early stop
        assert!(early_stop.check(0.8));

        // Improvement resets counter
        let mut early_stop2 = EarlyStopping::new(3, 1e-4);
        early_stop2.check(1.0);
        early_stop2.check(1.1); // no improvement
        early_stop2.check(1.2); // no improvement
        assert!(!early_stop2.check(0.5)); // improvement - resets
        assert_eq!(early_stop2.epochs_without_improvement(), 0);
    }

    #[test]
    fn test_lr_scheduler_reduce_on_plateau() {
        let mut scheduler = LRScheduler::new(
            LRSchedulerType::ReduceOnPlateau {
                factor: 0.5,
                patience: 2,
                min_lr: 1e-6,
            },
            0.01,
        );

        // Initial LR
        assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);

        // Good loss - no change
        scheduler.step(0, Some(1.0));
        assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);

        // No improvement for 2 epochs - triggers reduction
        scheduler.step(1, Some(1.1)); // worse
        scheduler.step(2, Some(1.2)); // worse again, triggers reduction
        assert!((scheduler.learning_rate() - 0.005).abs() < 1e-8);
    }

    #[test]
    fn test_lr_scheduler_cosine_annealing() {
        let mut scheduler = LRScheduler::new(
            LRSchedulerType::CosineAnnealing {
                t_max: 100,
                eta_min: 0.0001,
            },
            0.01,
        );

        // At epoch 0, LR should be at maximum (cos(0) = 1)
        // lr = eta_min + 0.5 * (initial_lr - eta_min) * (1 + 1) = initial_lr
        scheduler.step(0, None);
        assert!((scheduler.learning_rate() - 0.01).abs() < 1e-6);

        // At epoch t_max/2 (50), LR should be at midpoint (cos(π/2) = 0)
        // lr = eta_min + 0.5 * (initial_lr - eta_min) * (1 + 0) = (initial_lr + eta_min) / 2
        scheduler.step(50, None);
        let expected_mid = (0.01 + 0.0001) / 2.0; // 0.00505
        assert!(
            (scheduler.learning_rate() - expected_mid).abs() < 1e-6,
            "At midpoint: expected {}, got {}",
            expected_mid,
            scheduler.learning_rate()
        );

        // At epoch 99 (just before restart), LR should be near minimum
        scheduler.step(99, None);
        // t=99, cos(π*99/100) ≈ cos(0.99π) ≈ -0.9998
        // lr ≈ eta_min + 0.5 * 0.0099 * (1 - 0.9998) ≈ 0.0001
        assert!(
            scheduler.learning_rate() < 0.001,
            "Near end of cycle: expected near eta_min, got {}",
            scheduler.learning_rate()
        );

        // At epoch t_max (100), cycle restarts so LR goes back to max
        // Because epoch % t_max = 0
        scheduler.step(100, None);
        assert!(
            (scheduler.learning_rate() - 0.01).abs() < 1e-6,
            "Cycle restart: expected {}, got {}",
            0.01,
            scheduler.learning_rate()
        );
    }

    #[test]
    fn test_lr_scheduler_step_decay() {
        let mut scheduler = LRScheduler::new(
            LRSchedulerType::StepDecay {
                step_size: 10,
                gamma: 0.1,
            },
            0.01,
        );

        // Epochs 0-9: LR = 0.01
        scheduler.step(0, None);
        assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);
        scheduler.step(9, None);
        assert!((scheduler.learning_rate() - 0.01).abs() < 1e-8);

        // Epochs 10-19: LR = 0.001
        scheduler.step(10, None);
        assert!((scheduler.learning_rate() - 0.001).abs() < 1e-8);

        // Epochs 20-29: LR = 0.0001
        scheduler.step(20, None);
        assert!((scheduler.learning_rate() - 0.0001).abs() < 1e-8);
    }

    #[test]
    fn test_training_session_creation() {
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let session = TrainingSession::new(config, pde_config);

        assert!(!session.is_training());
    }

    // Integration test - runs actual training (slow)
    #[test]
    #[ignore = "slow integration test"]
    fn test_quick_training() {
        let config = TrainingConfig {
            epochs: 2,
            batch_size: 4,
            learning_rate: 0.001,
            n_train_samples: 8,
            n_val_samples: 4,
        };
        let pde_config = PDEConfig::darcy(16)
            .with_modes(4, 4)
            .with_width(8)
            .with_layers(2);

        // Explicitly use CpuBackend for test
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);
        let result = trainer.train();

        assert!(
            result.is_ok(),
            "Training should succeed: {:?}",
            result.err()
        );

        let final_progress = trainer.get_progress();
        assert_eq!(final_progress.status, TrainingStatus::Complete);
        assert_eq!(final_progress.epoch, 2);
    }

    #[test]
    fn test_trainer_save_weights_creates_file() {
        use tempfile::tempdir;

        let _temp_dir = tempdir().expect("Failed to create temp dir");
        let device = CpuDevice::default();

        let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
            .expect("Failed to create model");

        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);

        let weights_path = trainer
            .save_weights(&model)
            .expect("Failed to save weights");

        assert!(
            std::path::Path::new(&weights_path).exists(),
            "Weights file should exist"
        );
    }

    #[test]
    fn test_trainer_save_weights_roundtrip() {
        use rtx_neural_operator::weights::load_fno2d_weights;
        use tempfile::tempdir;

        let _temp_dir = tempdir().expect("Failed to create temp dir");
        let device = CpuDevice::default();

        let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
            .expect("Failed to create model");

        let original_weights = model.to_weights("test_pde");

        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device.clone());

        let weights_path = trainer
            .save_weights(&model)
            .expect("Failed to save weights");

        let loaded_weights = load_fno2d_weights(&weights_path).expect("Failed to load weights");

        assert_eq!(loaded_weights.config.width, original_weights.config.width);
        assert_eq!(
            loaded_weights.config.n_layers,
            original_weights.config.n_layers
        );
        assert_eq!(
            loaded_weights.config.n_modes,
            original_weights.config.n_modes
        );

        assert_eq!(
            loaded_weights.lifting_fc1_weight.len(),
            original_weights.lifting_fc1_weight.len()
        );
        for (a, b) in loaded_weights
            .lifting_fc1_weight
            .iter()
            .zip(original_weights.lifting_fc1_weight.iter())
        {
            assert!(
                (a - b).abs() < 1e-6,
                "Lifting fc1 weights mismatch: {} vs {}",
                a,
                b
            );
        }

        assert_eq!(
            loaded_weights.spectral_weights.len(),
            original_weights.spectral_weights.len()
        );
        for (loaded_sw, orig_sw) in loaded_weights
            .spectral_weights
            .iter()
            .zip(original_weights.spectral_weights.iter())
        {
            assert_eq!(loaded_sw.weights1_real.len(), orig_sw.weights1_real.len());
            for (a, b) in loaded_sw
                .weights1_real
                .iter()
                .zip(orig_sw.weights1_real.iter())
            {
                assert!(
                    (a - b).abs() < 1e-6,
                    "Spectral weights mismatch: {} vs {}",
                    a,
                    b
                );
            }
        }
    }

    #[test]
    fn test_trainer_save_weights_contains_expected_tensors() {
        use rtx_hub::safetensors::SafeTensors;
        use tempfile::tempdir;

        let _temp_dir = tempdir().expect("Failed to create temp dir");
        let device = CpuDevice::default();

        let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
            .expect("Failed to create model");

        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);

        let weights_path = trainer
            .save_weights(&model)
            .expect("Failed to save weights");

        let file_bytes = std::fs::read(&weights_path).expect("Failed to read weights file");
        let safetensors =
            SafeTensors::from_bytes(&file_bytes).expect("Failed to parse SafeTensors");

        let tensor_names = safetensors.tensor_names();

        assert!(
            tensor_names.contains(&"lifting.fcs.0.weight"),
            "Should contain lifting.fcs.0.weight"
        );
        assert!(
            tensor_names.contains(&"lifting.fcs.0.bias"),
            "Should contain lifting.fcs.0.bias"
        );
        assert!(
            tensor_names.contains(&"lifting.fcs.1.weight"),
            "Should contain lifting.fcs.1.weight"
        );
        assert!(
            tensor_names.contains(&"lifting.fcs.1.bias"),
            "Should contain lifting.fcs.1.bias"
        );

        assert!(
            tensor_names.contains(&"spectral_conv.0.weights1_real"),
            "Should contain spectral_conv.0.weights1_real"
        );
        assert!(
            tensor_names.contains(&"spectral_conv.0.weights1_imag"),
            "Should contain spectral_conv.0.weights1_imag"
        );
        assert!(
            tensor_names.contains(&"spectral_conv.1.weights1_real"),
            "Should contain spectral_conv.1.weights1_real"
        );

        assert!(
            tensor_names.contains(&"conv.0.weight"),
            "Should contain conv.0.weight"
        );
        assert!(
            tensor_names.contains(&"conv.0.bias"),
            "Should contain conv.0.bias"
        );

        assert!(
            tensor_names.contains(&"projection.0.weight"),
            "Should contain projection.0.weight"
        );
        assert!(
            tensor_names.contains(&"projection.0.bias"),
            "Should contain projection.0.bias"
        );
        assert!(
            tensor_names.contains(&"projection.1.weight"),
            "Should contain projection.1.weight"
        );
        assert!(
            tensor_names.contains(&"projection.1.bias"),
            "Should contain projection.1.bias"
        );
    }

    #[test]
    fn test_trainer_save_weights_includes_metadata() {
        use rtx_hub::safetensors::SafeTensors;
        use tempfile::tempdir;

        let _temp_dir = tempdir().expect("Failed to create temp dir");
        let device = CpuDevice::default();

        let model = FNO2d::<CpuBackend>::new_with_layers(1, 1, 8, 4, 2, &device)
            .expect("Failed to create model");

        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16)
            .with_modes(4, 4)
            .with_width(8)
            .with_layers(2);
        let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config.clone(), device);

        let weights_path = trainer
            .save_weights(&model)
            .expect("Failed to save weights");

        let file_bytes = std::fs::read(&weights_path).expect("Failed to read weights file");
        let safetensors =
            SafeTensors::from_bytes(&file_bytes).expect("Failed to parse SafeTensors");

        let metadata = safetensors.metadata().expect("Should have metadata");

        assert_eq!(
            metadata.get("model_type"),
            Some(&"FNO2d".to_string()),
            "Should have model_type metadata"
        );
        assert_eq!(
            metadata.get("pde_type"),
            Some(&"darcy_flow".to_string()),
            "Should have pde_type metadata"
        );
        assert_eq!(
            metadata.get("resolution"),
            Some(&"16".to_string()),
            "Should have resolution metadata"
        );
        assert_eq!(
            metadata.get("width"),
            Some(&"8".to_string()),
            "Should have width metadata"
        );
        assert_eq!(
            metadata.get("n_layers"),
            Some(&"2".to_string()),
            "Should have n_layers metadata"
        );
    }

    // ============================================================================
    // TDD Tests for GPU Acceleration (RED-GREEN-REFACTOR)
    // ============================================================================

    #[test]
    fn test_device_selection_reports_device_name() {
        // RED: Test that trainer reports device name in initial progress
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);

        let initial_progress = trainer.get_progress();

        // GREEN: Device name should be populated
        assert!(
            !initial_progress.device.is_empty(),
            "Device name should not be empty"
        );
        assert!(
            initial_progress.device.contains("CPU")
                || initial_progress.device.contains("CUDA")
                || initial_progress.device.contains("Metal"),
            "Device name should identify backend type: got '{}'",
            initial_progress.device
        );
    }

    #[test]
    fn test_training_progress_preserves_device_info() {
        // RED: Test that device info persists through training updates
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);

        let initial_device = trainer.get_progress().device.clone();

        // Update progress
        trainer.update_progress(|p| {
            p.epoch = 1;
            p.loss = 0.5;
        });

        let updated_progress = trainer.get_progress();

        // GREEN: Device name should be unchanged
        assert_eq!(
            updated_progress.device, initial_device,
            "Device name should persist through updates"
        );
        assert_eq!(updated_progress.epoch, 1);
    }

    #[test]
    fn test_cpu_backend_trainer_creation() {
        // RED: Test explicit CPU backend trainer creation
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let device = CpuDevice::default();

        let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);

        // GREEN: Verify trainer was created successfully
        assert!(!trainer.is_training());
        let progress = trainer.get_progress();
        assert!(
            progress.device.contains("CPU"),
            "Device should be CPU, got: {}",
            progress.device
        );
    }

    #[test]
    fn test_training_session_auto_selects_backend() {
        // RED: Test that TrainingSession automatically selects best backend
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);

        let session = TrainingSession::new(config, pde_config);

        // GREEN: Session should be created successfully with some backend
        assert!(!session.is_training());
        let progress = session.get_progress();
        assert!(!progress.device.is_empty(), "Device should be selected");
    }

    #[test]
    fn test_training_session_with_explicit_cpu_backend() {
        // RED: Test creating session with explicit CPU backend
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);

        let session = TrainingSession::with_backend::<CpuBackend>(config, pde_config);

        // GREEN: Should create session with CPU backend
        assert!(!session.is_training());
        let progress = session.get_progress();
        assert!(progress.device.contains("CPU"), "Should use CPU backend");
    }

    #[test]
    #[cfg(feature = "cuda")]
    fn test_training_session_with_cuda_backend() {
        // RED: Test creating session with CUDA backend if available
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);

        // Try to create CUDA device
        if CudaDevice::default_available() {
            let session = TrainingSession::with_backend::<CudaBackend>(config, pde_config);

            // GREEN: Should create session with CUDA backend
            assert!(!session.is_training());
            let progress = session.get_progress();
            assert!(
                progress.device.contains("CUDA"),
                "Should use CUDA backend, got: {}",
                progress.device
            );
        }
    }

    #[test]
    #[cfg(feature = "metal")]
    fn test_training_session_with_metal_backend() {
        // RED: Test creating session with Metal backend if available
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);

        // Try to create Metal device
        use rtx_backend_metal::MetalDevice;
        if MetalDevice::default_available() {
            let session = TrainingSession::with_backend::<MetalBackend>(config, pde_config);

            // GREEN: Should create session with Metal backend
            assert!(!session.is_training());
            let progress = session.get_progress();
            assert!(
                progress.device.contains("Metal"),
                "Should use Metal backend, got: {}",
                progress.device
            );
        }
    }

    #[test]
    fn test_get_device_string_returns_valid_name() {
        // RED: Test that get_device_string returns a valid device identifier
        let device_string = get_device_string();

        // GREEN: Should return non-empty string with backend type
        assert!(
            !device_string.is_empty(),
            "Device string should not be empty"
        );
        assert!(
            device_string.contains("CPU")
                || device_string.contains("CUDA")
                || device_string.contains("Metal")
                || device_string.contains("ROCm"),
            "Device string should contain backend type: got '{}'",
            device_string
        );
    }

    #[test]
    fn test_initial_progress_includes_learning_rate() {
        // RED: Test that initial progress includes the configured learning rate
        let config = TrainingConfig {
            epochs: 10,
            batch_size: 8,
            learning_rate: 0.002,
            n_train_samples: 100,
            n_val_samples: 20,
        };
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::new(config.clone(), pde_config);

        let initial_progress = trainer.get_progress();

        // GREEN: Initial LR should match config
        assert!(
            (initial_progress.current_lr - config.learning_rate).abs() < 1e-6,
            "Initial LR should be {}, got {}",
            config.learning_rate,
            initial_progress.current_lr
        );
    }

    #[test]
    fn test_backend_selection_order() {
        // RED: Test that select_best_backend returns available backend
        use rtx_backend::auto_select::select_best_backend;

        let backend = select_best_backend();

        // GREEN: Should return one of the valid backend types (case insensitive)
        let backend_str = backend.to_string();
        let backend_lower = backend_str.to_lowercase();
        assert!(
            backend_lower == "cpu"
                || backend_lower == "cuda"
                || backend_lower == "metal"
                || backend_lower == "rocm",
            "Backend selection returned unexpected type: {}",
            backend_str
        );
    }

    #[test]
    fn test_dynamic_trainer_trait_implementation() {
        // RED: Test that FnoTrainer implements DynamicTrainer
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let trainer = FnoTrainer::<CpuBackend>::new(config, pde_config);

        let dynamic: &dyn DynamicTrainer = &trainer;

        // GREEN: Should be able to call trait methods
        assert!(!dynamic.is_training());
        let progress = dynamic.get_progress();
        assert_eq!(progress.status, TrainingStatus::NotStarted);
    }

    #[test]
    fn test_training_session_creation_with_auto_backend() {
        // RED: Test that new() uses automatic backend selection
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);

        let session = TrainingSession::new(config, pde_config);
        let progress = session.get_progress();

        // GREEN: Should have selected some backend
        assert!(!progress.device.is_empty());

        // Device string should match one of the known formats
        let device_lower = progress.device.to_lowercase();
        assert!(
            device_lower.contains("cpu")
                || device_lower.contains("cuda")
                || device_lower.contains("metal")
                || device_lower.contains("rocm"),
            "Device should be recognized backend type: {}",
            progress.device
        );
    }

    #[test]
    fn test_progress_device_field_serialization() {
        // RED: Test that TrainingProgress with device field serializes correctly
        // Create a complete progress object with non-infinite values
        let mut progress = TrainingProgress::new(10, 5);
        progress.device = "CUDA (RTX 4090)".to_string();
        progress.current_lr = 0.001;
        progress.loss = 0.5;
        progress.best_loss = 0.3;
        progress.epoch = 5;
        progress.batch = 2;

        let json = serde_json::to_string(&progress).expect("Should serialize");
        let deserialized: TrainingProgress =
            serde_json::from_str(&json).expect("Should deserialize");

        // GREEN: Device field and learning rate should roundtrip correctly
        assert_eq!(deserialized.device, "CUDA (RTX 4090)");
        assert!((deserialized.current_lr - 0.001).abs() < 1e-6);
        assert!((deserialized.loss - 0.5).abs() < 1e-6);
        assert!((deserialized.best_loss - 0.3).abs() < 1e-6);
        assert_eq!(deserialized.epoch, 5);
        assert_eq!(deserialized.batch, 2);
    }

    #[test]
    fn test_trainer_reports_correct_device_for_cpu() {
        // RED: Test that CPU backend reports CPU device
        let config = TrainingConfig::quick();
        let pde_config = PDEConfig::darcy(16);
        let device = CpuDevice::default();
        let trainer = FnoTrainer::<CpuBackend>::with_device(config, pde_config, device);

        let progress = trainer.get_progress();

        // GREEN: Device should indicate CPU
        assert!(
            progress.device.contains("CPU"),
            "CPU backend should report CPU device, got: {}",
            progress.device
        );
    }

    #[test]
    fn test_multiple_trainers_with_different_backends() {
        // RED: Test creating multiple trainers with different backends
        let config1 = TrainingConfig::quick();
        let config2 = TrainingConfig::quick();
        let pde_config1 = PDEConfig::darcy(16);
        let pde_config2 = PDEConfig::darcy(16);

        let trainer1 = FnoTrainer::<CpuBackend>::new(config1, pde_config1);
        let session2 = TrainingSession::new(config2, pde_config2);

        // GREEN: Both should be created successfully
        assert!(!trainer1.is_training());
        assert!(!session2.is_training());

        let progress1 = trainer1.get_progress();
        let progress2 = session2.get_progress();

        assert!(!progress1.device.is_empty());
        assert!(!progress2.device.is_empty());
    }
}