rustytorch/demos/neural-operator-shared/src/ipc.rs

//! Inter-process communication types for Tauri frontend/backend
//!
//! This module defines all message types used for communication between
//! the Tauri frontend and the `RustyTorch`++ backend.
//!
//! # Example
//!
//! ```rust
//! use rtx_neural_operator_shared::config::PDEConfig;
//! use rtx_neural_operator_shared::ipc::{NeuralOperatorRequest, NeuralOperatorResponse};
//!
//! // Create initialization request
//! let config = PDEConfig::darcy(64);
//! let request = NeuralOperatorRequest::initialize(config);
//!
//! // Serialize for IPC
//! let json = serde_json::to_string(&request).unwrap();
//!
//! // Create success response
//! let response = NeuralOperatorResponse::initialized(64, 64);
//! ```

use serde::{Deserialize, Serialize};

use crate::config::PDEConfig;
use crate::error::NeuralOperatorError;

/// IPC request types from frontend to backend
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
pub enum NeuralOperatorRequest {
    /// Initialize the neural operator with configuration
    Initialize {
        /// PDE configuration
        config: PDEConfig,
    },

    /// Solve PDE with given input field
    Solve {
        /// Input field (flattened [H, W] array)
        input: Vec<f32>,
    },

    /// Set boundary conditions for the domain
    SetBoundaryConditions {
        /// Boundary mask (1.0 = boundary, 0.0 = interior)
        mask: Vec<f32>,
        /// Boundary values at masked points
        values: Vec<f32>,
    },

    /// Get current solution (cached from last solve)
    GetSolution,

    /// Get performance metrics
    GetMetrics,

    /// Reset to initial state
    Reset,

    /// Get model information
    GetModelInfo,
}

impl NeuralOperatorRequest {
    /// Creates an initialize request
    #[must_use]
    pub fn initialize(config: PDEConfig) -> Self {
        Self::Initialize { config }
    }

    /// Creates a solve request
    #[must_use]
    pub fn solve(input: Vec<f32>) -> Self {
        Self::Solve { input }
    }

    /// Creates a set boundary conditions request
    #[must_use]
    pub fn set_boundary_conditions(mask: Vec<f32>, values: Vec<f32>) -> Self {
        Self::SetBoundaryConditions { mask, values }
    }

    /// Creates a get solution request
    #[must_use]
    pub const fn get_solution() -> Self {
        Self::GetSolution
    }

    /// Creates a get metrics request
    #[must_use]
    pub const fn get_metrics() -> Self {
        Self::GetMetrics
    }

    /// Creates a reset request
    #[must_use]
    pub const fn reset() -> Self {
        Self::Reset
    }

    /// Creates a get model info request
    #[must_use]
    pub const fn get_model_info() -> Self {
        Self::GetModelInfo
    }
}

/// IPC response from backend to frontend
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum NeuralOperatorResponse {
    /// Model initialized successfully
    Initialized {
        /// Grid width
        width: u32,
        /// Grid height
        height: u32,
        /// Model name/description
        model_name: String,
    },

    /// Solution computed successfully
    Solution(SolutionData),

    /// Performance metrics
    Metrics(PerformanceMetrics),

    /// Model information
    ModelInfo(ModelInfo),

    /// Operation succeeded with no data
    Ok,

    /// Error occurred
    Error {
        /// Error code
        code: String,
        /// Error message
        message: String,
    },
}

impl NeuralOperatorResponse {
    /// Creates an initialized response
    #[must_use]
    pub fn initialized(width: u32, height: u32) -> Self {
        Self::Initialized {
            width,
            height,
            model_name: format!("FNO2d ({width}x{height})"),
        }
    }

    /// Creates a solution response
    #[must_use]
    pub fn solution(data: SolutionData) -> Self {
        Self::Solution(data)
    }

    /// Creates a metrics response
    #[must_use]
    pub fn metrics(metrics: PerformanceMetrics) -> Self {
        Self::Metrics(metrics)
    }

    /// Creates a model info response
    #[must_use]
    pub fn model_info(info: ModelInfo) -> Self {
        Self::ModelInfo(info)
    }

    /// Creates an OK response
    #[must_use]
    pub const fn ok() -> Self {
        Self::Ok
    }

    /// Creates an error response from an error
    #[must_use]
    pub fn error(err: NeuralOperatorError) -> Self {
        Self::Error {
            code: err.code().to_string(),
            message: err.to_string(),
        }
    }

    /// Creates an error response from a message
    #[must_use]
    pub fn error_message(message: impl Into<String>) -> Self {
        Self::Error {
            code: "ERROR".to_string(),
            message: message.into(),
        }
    }

    /// Returns whether the response indicates success
    #[must_use]
    pub const fn is_success(&self) -> bool {
        !matches!(self, Self::Error { .. })
    }
}

/// Solution data from neural operator inference
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct SolutionData {
    /// Solution field (flattened [H, W] array)
    pub solution: Vec<f32>,

    /// Grid width
    pub width: u32,

    /// Grid height
    pub height: u32,

    /// Neural operator inference time in milliseconds
    pub inference_time_ms: f64,

    /// FEM baseline time in milliseconds (if computed)
    pub fem_time_ms: Option<f64>,

    /// Minimum value in solution
    pub min_value: f32,

    /// Maximum value in solution
    pub max_value: f32,

    /// Mean value in solution
    pub mean_value: f32,
}

impl SolutionData {
    /// Creates new solution data
    #[must_use]
    pub fn new(solution: Vec<f32>, width: u32, height: u32, inference_time_ms: f64) -> Self {
        let (min_value, max_value, sum) = solution.iter().fold(
            (f32::INFINITY, f32::NEG_INFINITY, 0.0_f64),
            |(min, max, sum), &v| (min.min(v), max.max(v), sum + f64::from(v)),
        );
        let mean_value = (sum / solution.len() as f64) as f32;

        Self {
            solution,
            width,
            height,
            inference_time_ms,
            fem_time_ms: None,
            min_value,
            max_value,
            mean_value,
        }
    }

    /// Adds FEM baseline timing
    #[must_use]
    pub fn with_fem_time(mut self, fem_time_ms: f64) -> Self {
        self.fem_time_ms = Some(fem_time_ms);
        self
    }

    /// Returns the speedup factor vs FEM (if available)
    #[must_use]
    pub fn speedup_factor(&self) -> Option<f64> {
        self.fem_time_ms.map(|fem| fem / self.inference_time_ms)
    }
}

/// Performance metrics for monitoring
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct PerformanceMetrics {
    /// Average inference time in milliseconds
    pub avg_inference_time_ms: f64,

    /// Minimum inference time in milliseconds
    pub min_inference_time_ms: f64,

    /// Maximum inference time in milliseconds
    pub max_inference_time_ms: f64,

    /// Number of inferences performed
    pub inference_count: u64,

    /// Memory usage in megabytes
    pub memory_usage_mb: f32,

    /// Throughput (inferences per second)
    pub throughput: f64,
}

impl PerformanceMetrics {
    /// Creates new performance metrics
    #[must_use]
    pub fn new() -> Self {
        Self {
            avg_inference_time_ms: 0.0,
            min_inference_time_ms: f64::INFINITY,
            max_inference_time_ms: 0.0,
            inference_count: 0,
            memory_usage_mb: 0.0,
            throughput: 0.0,
        }
    }

    /// Records a new inference timing
    pub fn record_inference(&mut self, time_ms: f64) {
        self.inference_count += 1;
        self.min_inference_time_ms = self.min_inference_time_ms.min(time_ms);
        self.max_inference_time_ms = self.max_inference_time_ms.max(time_ms);

        // Running average
        let n = self.inference_count as f64;
        self.avg_inference_time_ms = self.avg_inference_time_ms * (n - 1.0) / n + time_ms / n;

        // Throughput based on average
        if self.avg_inference_time_ms > 0.0 {
            self.throughput = 1000.0 / self.avg_inference_time_ms;
        }
    }

    /// Sets memory usage
    #[must_use]
    pub const fn with_memory(mut self, memory_mb: f32) -> Self {
        self.memory_usage_mb = memory_mb;
        self
    }

    /// Returns whether performance meets real-time requirements (<33ms for 30fps)
    #[must_use]
    pub fn is_realtime(&self) -> bool {
        self.avg_inference_time_ms < 33.0
    }
}

impl Default for PerformanceMetrics {
    fn default() -> Self {
        Self::new()
    }
}

/// Model information
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Model name
    pub name: String,

    /// PDE type
    pub pde_type: String,

    /// Input resolution
    pub resolution: u32,

    /// Number of Fourier modes
    pub n_modes: (u32, u32),

    /// Model width (hidden dimension)
    pub model_width: u32,

    /// Number of layers
    pub n_layers: u32,

    /// Total parameters
    pub total_params: u64,

    /// Model size in megabytes
    pub size_mb: f32,
}

impl ModelInfo {
    /// Creates new model info
    #[must_use]
    pub fn new(
        name: impl Into<String>,
        pde_type: impl Into<String>,
        resolution: u32,
        n_modes: (u32, u32),
        model_width: u32,
        n_layers: u32,
    ) -> Self {
        // Rough parameter count estimate for FNO2d
        let lifting_params = (3 * 2 * model_width) + (2 * model_width * model_width);
        let spectral_params = n_layers * 2 * model_width * model_width * n_modes.0 * n_modes.1 * 2;
        let conv_params = n_layers * (model_width * model_width + model_width);
        let projection_params = (model_width * 128) + 128;
        let total_params =
            u64::from(lifting_params + spectral_params + conv_params + projection_params);

        Self {
            name: name.into(),
            pde_type: pde_type.into(),
            resolution,
            n_modes,
            model_width,
            n_layers,
            total_params,
            size_mb: (total_params * 4) as f32 / (1024.0 * 1024.0), // 4 bytes per f32
        }
    }
}

/// Helper function to get current timestamp in milliseconds
#[must_use]
pub fn current_timestamp_ms() -> u64 {
    use std::time::{SystemTime, UNIX_EPOCH};
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

// =============================================================================
// BENCHMARK TYPES
// =============================================================================

/// Benchmark request from frontend to backend
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "payload")]
pub enum BenchmarkRequest {
    /// Run benchmark comparing solvers
    RunBenchmark {
        /// Resolutions to test (e.g., [32, 64, 128, 256])
        resolutions: Vec<usize>,
        /// Methods to benchmark (e.g., ["FNO", "FDM", "FEM"])
        methods: Vec<String>,
        /// PDE type ("Poisson", "Heat", "Darcy")
        pde_type: String,
        /// Number of trials for averaging
        n_trials: usize,
    },

    /// Get status of running benchmark
    GetBenchmarkStatus,

    /// Cancel running benchmark
    CancelBenchmark,
}

impl BenchmarkRequest {
    /// Creates a run benchmark request
    #[must_use]
    pub fn run(
        resolutions: Vec<usize>,
        methods: Vec<String>,
        pde_type: String,
        n_trials: usize,
    ) -> Self {
        Self::RunBenchmark {
            resolutions,
            methods,
            pde_type,
            n_trials,
        }
    }

    /// Creates a get status request
    #[must_use]
    pub const fn get_status() -> Self {
        Self::GetBenchmarkStatus
    }

    /// Creates a cancel request
    #[must_use]
    pub const fn cancel() -> Self {
        Self::CancelBenchmark
    }
}

/// Benchmark result data (serializable version of `BenchmarkResult`)
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BenchmarkResultData {
    /// Solver method name
    pub method: String,
    /// Grid resolution
    pub resolution: usize,
    /// Solve time in milliseconds
    pub solve_time_ms: f64,
    /// L2 error vs reference (if available)
    pub l2_error: Option<f64>,
    /// Maximum error vs reference (if available)
    pub max_error: Option<f64>,
    /// Memory usage in megabytes
    pub memory_mb: f64,
    /// Number of iterations (for iterative methods)
    pub iterations: Option<usize>,
    /// PDE type
    pub pde_type: String,
}

impl BenchmarkResultData {
    /// Creates a new benchmark result
    #[must_use]
    pub fn new(method: impl Into<String>, resolution: usize, pde_type: impl Into<String>) -> Self {
        Self {
            method: method.into(),
            resolution,
            solve_time_ms: 0.0,
            l2_error: None,
            max_error: None,
            memory_mb: 0.0,
            iterations: None,
            pde_type: pde_type.into(),
        }
    }

    /// Sets solve time
    #[must_use]
    pub const fn with_time(mut self, time_ms: f64) -> Self {
        self.solve_time_ms = time_ms;
        self
    }

    /// Sets L2 error
    #[must_use]
    pub const fn with_l2_error(mut self, error: f64) -> Self {
        self.l2_error = Some(error);
        self
    }

    /// Sets max error
    #[must_use]
    pub const fn with_max_error(mut self, error: f64) -> Self {
        self.max_error = Some(error);
        self
    }

    /// Sets memory usage
    #[must_use]
    pub const fn with_memory(mut self, memory_mb: f64) -> Self {
        self.memory_mb = memory_mb;
        self
    }

    /// Sets iterations
    #[must_use]
    pub const fn with_iterations(mut self, iterations: usize) -> Self {
        self.iterations = Some(iterations);
        self
    }

    /// Computes speedup vs a baseline time
    #[must_use]
    pub fn speedup_vs(&self, baseline_time_ms: f64) -> f64 {
        baseline_time_ms / self.solve_time_ms
    }
}

/// Summary statistics for a solver across multiple runs (serializable)
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct BenchmarkSummaryData {
    /// Solver method name
    pub method: String,
    /// Resolution
    pub resolution: usize,
    /// Average solve time in milliseconds
    pub avg_time_ms: f64,
    /// Standard deviation of solve time
    pub std_time_ms: f64,
    /// Minimum solve time
    pub min_time_ms: f64,
    /// Maximum solve time
    pub max_time_ms: f64,
    /// Average L2 error (if available)
    pub avg_l2_error: Option<f64>,
    /// Average memory usage
    pub avg_memory_mb: f64,
    /// Number of runs
    pub n_runs: usize,
    /// PDE type
    pub pde_type: String,
}

impl BenchmarkSummaryData {
    /// Computes speedup factor vs another solver
    #[must_use]
    pub fn speedup_vs(&self, other: &Self) -> f64 {
        other.avg_time_ms / self.avg_time_ms
    }
}

/// Benchmark response from backend to frontend
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", content = "data")]
pub enum BenchmarkResponse {
    /// Benchmark results computed
    Results {
        /// Individual benchmark results
        results: Vec<BenchmarkResultData>,
        /// Summary statistics (if multiple trials)
        summaries: Option<Vec<BenchmarkSummaryData>>,
        /// Speedup analysis report
        speedup_report: Option<String>,
    },

    /// Benchmark status update
    Status {
        /// Current progress (0.0 to 1.0)
        progress: f64,
        /// Status message
        message: String,
        /// Current method being benchmarked
        current_method: Option<String>,
        /// Current resolution being tested
        current_resolution: Option<usize>,
    },

    /// Benchmark completed
    Complete {
        /// Final results
        results: Vec<BenchmarkResultData>,
    },

    /// Benchmark cancelled
    Cancelled,

    /// Error occurred
    Error {
        /// Error code
        code: String,
        /// Error message
        message: String,
    },
}

impl BenchmarkResponse {
    /// Creates a results response
    #[must_use]
    pub fn results(results: Vec<BenchmarkResultData>) -> Self {
        Self::Results {
            results,
            summaries: None,
            speedup_report: None,
        }
    }

    /// Creates a results response with summaries
    #[must_use]
    pub fn results_with_summaries(
        results: Vec<BenchmarkResultData>,
        summaries: Vec<BenchmarkSummaryData>,
    ) -> Self {
        Self::Results {
            results,
            summaries: Some(summaries),
            speedup_report: None,
        }
    }

    /// Creates a results response with summaries and speedup report
    #[must_use]
    pub fn results_with_analysis(
        results: Vec<BenchmarkResultData>,
        summaries: Vec<BenchmarkSummaryData>,
        speedup_report: String,
    ) -> Self {
        Self::Results {
            results,
            summaries: Some(summaries),
            speedup_report: Some(speedup_report),
        }
    }

    /// Creates a status response
    #[must_use]
    pub fn status(progress: f64, message: impl Into<String>) -> Self {
        Self::Status {
            progress,
            message: message.into(),
            current_method: None,
            current_resolution: None,
        }
    }

    /// Creates a complete response
    #[must_use]
    pub fn complete(results: Vec<BenchmarkResultData>) -> Self {
        Self::Complete { results }
    }

    /// Creates a cancelled response
    #[must_use]
    pub const fn cancelled() -> Self {
        Self::Cancelled
    }

    /// Creates an error response
    #[must_use]
    pub fn error(code: impl Into<String>, message: impl Into<String>) -> Self {
        Self::Error {
            code: code.into(),
            message: message.into(),
        }
    }

    /// Returns whether the response indicates success
    #[must_use]
    pub const fn is_success(&self) -> bool {
        !matches!(self, Self::Error { .. })
    }
}

// =============================================================================
// TRAINING TYPES
// =============================================================================

/// Training configuration from frontend
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainingConfig {
    /// Number of training epochs
    pub epochs: usize,

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

    /// Learning rate
    pub learning_rate: f32,

    /// Number of training samples to generate
    pub n_train_samples: usize,

    /// Number of validation samples
    pub n_val_samples: usize,
}

impl Default for TrainingConfig {
    fn default() -> Self {
        Self {
            epochs: 50,
            batch_size: 16,
            learning_rate: 1e-3,
            n_train_samples: 1000,
            n_val_samples: 100,
        }
    }
}

impl TrainingConfig {
    /// Creates a new training config with default values
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Sets the number of epochs
    #[must_use]
    pub const fn with_epochs(mut self, epochs: usize) -> Self {
        self.epochs = epochs;
        self
    }

    /// Sets the batch size
    #[must_use]
    pub const fn with_batch_size(mut self, batch_size: usize) -> Self {
        self.batch_size = batch_size;
        self
    }

    /// Sets the learning rate
    #[must_use]
    pub const fn with_learning_rate(mut self, learning_rate: f32) -> Self {
        self.learning_rate = learning_rate;
        self
    }

    /// Sets the number of training samples
    #[must_use]
    pub const fn with_train_samples(mut self, n_samples: usize) -> Self {
        self.n_train_samples = n_samples;
        self
    }

    /// Quick training preset (for testing)
    #[must_use]
    pub fn quick() -> Self {
        Self {
            epochs: 10,
            batch_size: 8,
            learning_rate: 1e-3,
            n_train_samples: 100,
            n_val_samples: 20,
        }
    }

    /// Standard training preset
    #[must_use]
    pub fn standard() -> Self {
        Self::default()
    }

    /// Extended training preset (higher quality)
    #[must_use]
    pub fn extended() -> Self {
        Self {
            epochs: 200,
            batch_size: 32,
            learning_rate: 5e-4,
            n_train_samples: 5000,
            n_val_samples: 500,
        }
    }
}

/// Training status enum
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(tag = "status", content = "details")]
pub enum TrainingStatus {
    /// Training has not started
    NotStarted,

    /// Generating training data
    GeneratingData {
        /// Samples generated so far
        samples_generated: usize,
        /// Total samples to generate
        total_samples: usize,
    },

    /// Training in progress
    Training,

    /// Training completed successfully
    Complete,

    /// Training was cancelled by user
    Cancelled,

    /// Training failed with error
    Error(String),
}

impl TrainingStatus {
    /// Returns whether training is in progress
    #[must_use]
    pub fn is_running(&self) -> bool {
        matches!(self, Self::GeneratingData { .. } | Self::Training)
    }

    /// Returns whether training has finished (success, cancel, or error)
    #[must_use]
    pub fn is_finished(&self) -> bool {
        matches!(self, Self::Complete | Self::Cancelled | Self::Error(_))
    }
}

/// Training progress update (sent to frontend via polling)
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TrainingProgress {
    /// Current epoch (1-indexed)
    pub epoch: usize,

    /// Total number of epochs
    pub total_epochs: usize,

    /// Current batch within epoch
    pub batch: usize,

    /// Total batches per epoch
    pub total_batches: usize,

    /// Current training loss
    pub loss: f32,

    /// Best loss achieved so far
    pub best_loss: f32,

    /// Validation loss (if available)
    pub val_loss: Option<f32>,

    /// Training samples processed per second
    pub samples_per_sec: f32,

    /// Estimated time remaining in seconds
    pub eta_seconds: f32,

    /// Elapsed time in seconds
    pub elapsed_seconds: f32,

    /// Current training status
    pub status: TrainingStatus,

    /// Loss history (last N epochs)
    pub loss_history: Vec<f32>,

    /// Device used for training (e.g., "CPU", "CUDA (RTX 4090)", "Metal (M2 Max)")
    #[serde(default)]
    pub device: String,

    /// Current learning rate (may change with LR scheduling)
    #[serde(default)]
    pub current_lr: f32,
}

impl TrainingProgress {
    /// Creates a new training progress with initial values
    #[must_use]
    pub fn new(total_epochs: usize, total_batches: usize) -> Self {
        Self {
            epoch: 0,
            total_epochs,
            batch: 0,
            total_batches,
            loss: f32::INFINITY,
            best_loss: f32::INFINITY,
            val_loss: None,
            samples_per_sec: 0.0,
            eta_seconds: 0.0,
            elapsed_seconds: 0.0,
            status: TrainingStatus::NotStarted,
            loss_history: Vec::new(),
            device: String::from("CPU"),
            current_lr: 0.0,
        }
    }

    /// Creates progress indicating data generation phase
    #[must_use]
    pub fn generating_data(samples_generated: usize, total_samples: usize) -> Self {
        Self {
            epoch: 0,
            total_epochs: 0,
            batch: 0,
            total_batches: 0,
            loss: f32::INFINITY,
            best_loss: f32::INFINITY,
            val_loss: None,
            samples_per_sec: 0.0,
            eta_seconds: 0.0,
            elapsed_seconds: 0.0,
            status: TrainingStatus::GeneratingData {
                samples_generated,
                total_samples,
            },
            loss_history: Vec::new(),
            device: String::from("CPU"),
            current_lr: 0.0,
        }
    }

    /// Creates progress indicating completion
    #[must_use]
    pub fn complete(best_loss: f32, elapsed_seconds: f32, loss_history: Vec<f32>) -> Self {
        Self {
            epoch: loss_history.len(),
            total_epochs: loss_history.len(),
            batch: 0,
            total_batches: 0,
            loss: best_loss,
            best_loss,
            val_loss: None,
            samples_per_sec: 0.0,
            eta_seconds: 0.0,
            elapsed_seconds,
            status: TrainingStatus::Complete,
            loss_history,
            device: String::from("CPU"),
            current_lr: 0.0,
        }
    }

    /// Creates progress indicating cancellation
    #[must_use]
    pub fn cancelled() -> Self {
        Self {
            status: TrainingStatus::Cancelled,
            ..Self::new(0, 0)
        }
    }

    /// Creates progress indicating error
    #[must_use]
    pub fn error(message: impl Into<String>) -> Self {
        Self {
            status: TrainingStatus::Error(message.into()),
            ..Self::new(0, 0)
        }
    }

    /// Sets the device name for training
    #[must_use]
    pub fn with_device(mut self, device: impl Into<String>) -> Self {
        self.device = device.into();
        self
    }

    /// Sets the current learning rate
    #[must_use]
    pub fn with_learning_rate(mut self, lr: f32) -> Self {
        self.current_lr = lr;
        self
    }

    /// Returns the progress as a percentage (0.0 to 1.0)
    #[must_use]
    pub fn progress_fraction(&self) -> f32 {
        if self.total_epochs == 0 {
            return 0.0;
        }
        let epoch_progress = self.epoch as f32 / self.total_epochs as f32;
        let batch_progress = if self.total_batches > 0 {
            self.batch as f32 / self.total_batches as f32 / self.total_epochs as f32
        } else {
            0.0
        };
        epoch_progress + batch_progress
    }
}

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

    #[test]
    fn test_request_serialization() {
        let config = PDEConfig::darcy(64);
        let request = NeuralOperatorRequest::initialize(config);

        let json = serde_json::to_string(&request).unwrap();
        let deserialized: NeuralOperatorRequest = serde_json::from_str(&json).unwrap();

        assert!(matches!(
            deserialized,
            NeuralOperatorRequest::Initialize { .. }
        ));
    }

    #[test]
    fn test_response_success() {
        let response = NeuralOperatorResponse::initialized(64, 64);
        assert!(response.is_success());
    }

    #[test]
    fn test_response_error() {
        let response = NeuralOperatorResponse::error(NeuralOperatorError::NotInitialized);
        assert!(!response.is_success());
    }

    #[test]
    fn test_solution_data_stats() {
        let solution = vec![1.0, 2.0, 3.0, 4.0];
        let data = SolutionData::new(solution, 2, 2, 10.0);

        assert!((data.min_value - 1.0).abs() < f32::EPSILON);
        assert!((data.max_value - 4.0).abs() < f32::EPSILON);
        assert!((data.mean_value - 2.5).abs() < f32::EPSILON);
    }

    #[test]
    fn test_speedup_factor() {
        let solution = vec![1.0; 4];
        let data = SolutionData::new(solution, 2, 2, 10.0).with_fem_time(1000.0);

        assert!((data.speedup_factor().unwrap() - 100.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_performance_metrics_recording() {
        let mut metrics = PerformanceMetrics::new();
        metrics.record_inference(10.0);
        metrics.record_inference(20.0);

        assert_eq!(metrics.inference_count, 2);
        assert!((metrics.avg_inference_time_ms - 15.0).abs() < f64::EPSILON);
        assert!((metrics.min_inference_time_ms - 10.0).abs() < f64::EPSILON);
        assert!((metrics.max_inference_time_ms - 20.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_model_info() {
        let info = ModelInfo::new("FNO2d", "darcy_flow", 64, (12, 12), 32, 4);

        assert!(info.total_params > 0);
        assert!(info.size_mb > 0.0);
    }

    // =============================================================================
    // BENCHMARK IPC TESTS (RED PHASE - THESE WILL FAIL UNTIL WE IMPLEMENT)
    // =============================================================================

    #[test]
    fn test_benchmark_request_creation() {
        let request = BenchmarkRequest::run(
            vec![32, 64, 128],
            vec!["FNO".to_string(), "FDM".to_string(), "FEM".to_string()],
            "Poisson".to_string(),
            5,
        );

        assert!(matches!(request, BenchmarkRequest::RunBenchmark { .. }));
    }

    #[test]
    fn test_benchmark_request_serialization() {
        let request =
            BenchmarkRequest::run(vec![64], vec!["FNO".to_string()], "Darcy".to_string(), 1);

        let json = serde_json::to_string(&request).unwrap();
        let deserialized: BenchmarkRequest = serde_json::from_str(&json).unwrap();

        assert_eq!(request, deserialized);
    }

    #[test]
    fn test_benchmark_result_data_creation() {
        let result = BenchmarkResultData {
            method: "FNO".to_string(),
            resolution: 64,
            solve_time_ms: 5.0,
            l2_error: Some(0.001),
            max_error: Some(0.01),
            memory_mb: 10.5,
            iterations: None,
            pde_type: "Poisson".to_string(),
        };

        assert_eq!(result.method, "FNO");
        assert_eq!(result.resolution, 64);
        assert!(result.solve_time_ms > 0.0);
    }

    #[test]
    fn test_benchmark_response_creation() {
        let results = vec![
            BenchmarkResultData {
                method: "FNO".to_string(),
                resolution: 64,
                solve_time_ms: 5.0,
                l2_error: Some(0.001),
                max_error: None,
                memory_mb: 10.0,
                iterations: None,
                pde_type: "Poisson".to_string(),
            },
            BenchmarkResultData {
                method: "FDM-SOR".to_string(),
                resolution: 64,
                solve_time_ms: 500.0,
                l2_error: Some(0.0001),
                max_error: None,
                memory_mb: 2.0,
                iterations: Some(1500),
                pde_type: "Poisson".to_string(),
            },
        ];

        let response = BenchmarkResponse::results(results);
        assert!(matches!(response, BenchmarkResponse::Results { .. }));
    }

    #[test]
    fn test_benchmark_response_serialization() {
        let results = vec![BenchmarkResultData {
            method: "FEM-CG".to_string(),
            resolution: 32,
            solve_time_ms: 100.0,
            l2_error: Some(0.0005),
            max_error: Some(0.005),
            memory_mb: 3.5,
            iterations: Some(800),
            pde_type: "Heat".to_string(),
        }];

        let response = BenchmarkResponse::results(results);
        let json = serde_json::to_string(&response).unwrap();
        let deserialized: BenchmarkResponse = serde_json::from_str(&json).unwrap();

        assert!(matches!(deserialized, BenchmarkResponse::Results { .. }));
    }

    #[test]
    fn test_benchmark_speedup_computation() {
        let result_data = BenchmarkResultData {
            method: "FNO".to_string(),
            resolution: 128,
            solve_time_ms: 10.0,
            l2_error: Some(0.002),
            max_error: None,
            memory_mb: 15.0,
            iterations: None,
            pde_type: "Darcy".to_string(),
        };

        let baseline_time = 1000.0;
        let speedup = result_data.speedup_vs(baseline_time);
        assert!((speedup - 100.0).abs() < f64::EPSILON);
    }

    #[test]
    fn test_benchmark_summary_data() {
        let summary = BenchmarkSummaryData {
            method: "FDM-SOR".to_string(),
            resolution: 64,
            avg_time_ms: 450.0,
            std_time_ms: 25.0,
            min_time_ms: 420.0,
            max_time_ms: 480.0,
            avg_l2_error: Some(0.0001),
            avg_memory_mb: 2.0,
            n_runs: 10,
            pde_type: "Poisson".to_string(),
        };

        assert_eq!(summary.n_runs, 10);
        assert!(summary.avg_time_ms > summary.min_time_ms);
        assert!(summary.avg_time_ms < summary.max_time_ms);
    }
}