rustytorch/demos/rtx-neural-operator-demo/src/benchmark.rs

//! PDE Solver Benchmark Module
//!
//! Provides benchmark comparisons between:
//! - FNO (Fourier Neural Operator) - learned operator
//! - FEM (Finite Element Method) - classical numerical method
//! - FDM (Finite Difference Method) - classical numerical method
//!
//! # Example
//!
//! ```rust,ignore
//! use rtx_neural_operator_demo::benchmark::{BenchmarkConfig, BenchmarkRunner};
//!
//! let config = BenchmarkConfig::default();
//! let runner = BenchmarkRunner::new(config);
//! let results = runner.run_benchmark()?;
//!
//! for result in &results {
//!     println!("{}: {:.2}ms, L2 error: {:.6}", result.method, result.solve_time_ms, result.l2_error);
//! }
//! ```

use std::time::Instant;

use rtx_backend_cpu::{CpuBackend, CpuDevice};
use rtx_neural_operator::FNO2d;
use rtx_nn::GenericModule4D;
use rtx_tensor::GenericTensor;

// ============================================================================
// Benchmark Configuration
// ============================================================================

/// Configuration for benchmark runs
#[derive(Debug, Clone)]
pub struct BenchmarkConfig {
    /// Resolutions to test (e.g., [32, 64, 128, 256])
    pub resolutions: Vec<usize>,
    /// Number of test problems per resolution
    pub n_problems: usize,
    /// PDE type to benchmark
    pub pde_type: PDEType,
    /// Whether to use high-resolution reference for error computation
    pub use_reference: bool,
    /// Reference resolution for computing ground truth
    pub reference_resolution: usize,
}

impl Default for BenchmarkConfig {
    fn default() -> Self {
        Self {
            resolutions: vec![32, 64, 128],
            n_problems: 10,
            pde_type: PDEType::Poisson,
            use_reference: true,
            reference_resolution: 512,
        }
    }
}

impl BenchmarkConfig {
    /// Create config for quick benchmarking (fewer problems, lower resolutions)
    #[must_use]
    pub fn quick() -> Self {
        Self {
            resolutions: vec![32, 64],
            n_problems: 3,
            pde_type: PDEType::Poisson,
            use_reference: false,
            reference_resolution: 256,
        }
    }

    /// Create config for comprehensive benchmarking
    #[must_use]
    pub fn comprehensive() -> Self {
        Self {
            resolutions: vec![32, 64, 128, 256],
            n_problems: 20,
            pde_type: PDEType::Poisson,
            use_reference: true,
            reference_resolution: 512,
        }
    }

    /// Set PDE type
    #[must_use]
    pub fn with_pde_type(mut self, pde_type: PDEType) -> Self {
        self.pde_type = pde_type;
        self
    }
}

/// PDE types for benchmarking
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PDEType {
    /// Poisson equation: -∇²u = f
    Poisson,
    /// Heat equation (steady state): -∇·(k∇u) = f
    Heat,
    /// Darcy flow: -∇·(a∇u) = f with variable coefficient
    Darcy,
}

impl PDEType {
    /// Get human-readable name
    #[must_use]
    pub fn name(&self) -> &'static str {
        match self {
            Self::Poisson => "Poisson",
            Self::Heat => "Heat",
            Self::Darcy => "Darcy",
        }
    }
}

// ============================================================================
// Benchmark Results
// ============================================================================

/// Result from a single benchmark run
#[derive(Debug, Clone)]
pub struct BenchmarkResult {
    /// 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 (estimated)
    pub memory_mb: f64,
    /// Number of iterations (for iterative methods)
    pub iterations: Option<usize>,
    /// PDE type
    pub pde_type: PDEType,
}

impl BenchmarkResult {
    /// Create a new benchmark result
    pub fn new(method: impl Into<String>, resolution: usize, pde_type: PDEType) -> 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,
        }
    }

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

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

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

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

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

/// Summary statistics for a solver across multiple runs
#[derive(Debug, Clone)]
pub struct BenchmarkSummary {
    /// 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,
}

impl BenchmarkSummary {
    /// Compute summary from a list of results
    pub fn from_results(results: &[BenchmarkResult]) -> Option<Self> {
        if results.is_empty() {
            return None;
        }

        let method = results[0].method.clone();
        let resolution = results[0].resolution;
        let n_runs = results.len();

        let times: Vec<f64> = results.iter().map(|r| r.solve_time_ms).collect();
        let avg_time_ms = times.iter().sum::<f64>() / n_runs as f64;
        let min_time_ms = times.iter().copied().fold(f64::INFINITY, f64::min);
        let max_time_ms = times.iter().copied().fold(f64::NEG_INFINITY, f64::max);

        let variance = times.iter().map(|t| (t - avg_time_ms).powi(2)).sum::<f64>() / n_runs as f64;
        let std_time_ms = variance.sqrt();

        let avg_l2_error = if results.iter().all(|r| r.l2_error.is_some()) {
            Some(results.iter().filter_map(|r| r.l2_error).sum::<f64>() / n_runs as f64)
        } else {
            None
        };

        let avg_memory_mb = results.iter().map(|r| r.memory_mb).sum::<f64>() / n_runs as f64;

        Some(Self {
            method,
            resolution,
            avg_time_ms,
            std_time_ms,
            min_time_ms,
            max_time_ms,
            avg_l2_error,
            avg_memory_mb,
            n_runs,
        })
    }

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

// ============================================================================
// Finite Difference Method (FDM) Solver
// ============================================================================

/// FDM solver using Jacobi iteration for 2D Poisson equation.
///
/// Solves: -∇²u = f on [0,1]² with Dirichlet boundary conditions u = 0.
///
/// Uses the 5-point stencil:
/// ```text
///        u[i,j+1]
///           |
/// u[i-1,j] - u[i,j] - u[i+1,j]
///           |
///        u[i,j-1]
///
/// -∇²u ≈ (4*u[i,j] - u[i-1,j] - u[i+1,j] - u[i,j-1] - u[i,j+1]) / h²
/// ```
pub struct FdmSolver {
    /// Grid resolution (N x N)
    resolution: usize,
    /// Grid spacing h = 1/(N-1)
    h: f64,
    /// Maximum iterations
    max_iterations: usize,
    /// Convergence tolerance
    tolerance: f64,
}

impl FdmSolver {
    /// Create a new FDM solver
    #[must_use]
    pub fn new(resolution: usize) -> Self {
        let h = 1.0 / (resolution - 1) as f64;
        Self {
            resolution,
            h,
            max_iterations: 10000,
            tolerance: 1e-6,
        }
    }

    /// Set maximum iterations
    #[must_use]
    pub fn with_max_iterations(mut self, max_iter: usize) -> Self {
        self.max_iterations = max_iter;
        self
    }

    /// Set convergence tolerance
    #[must_use]
    pub fn with_tolerance(mut self, tol: f64) -> Self {
        self.tolerance = tol;
        self
    }

    /// Solve -∇²u = f with zero Dirichlet boundary conditions using Jacobi iteration.
    ///
    /// # Arguments
    /// * `rhs` - Right-hand side f (flattened N x N array)
    ///
    /// # Returns
    /// Tuple of (solution, iterations)
    #[must_use]
    pub fn solve_jacobi(&self, rhs: &[f64]) -> (Vec<f64>, usize) {
        let n = self.resolution;
        let h2 = self.h * self.h;

        // Initialize solution to zero
        let mut u = vec![0.0; n * n];
        let mut u_new = vec![0.0; n * n];

        let mut iterations = 0;

        for iter in 0..self.max_iterations {
            let mut max_diff = 0.0f64;

            // Update interior points only
            for i in 1..n - 1 {
                for j in 1..n - 1 {
                    let idx = i * n + j;
                    let left = u[(i - 1) * n + j];
                    let right = u[(i + 1) * n + j];
                    let down = u[i * n + (j - 1)];
                    let up = u[i * n + (j + 1)];

                    // Jacobi update: u_new = (u_left + u_right + u_down + u_up + h²*f) / 4
                    u_new[idx] = (left + right + down + up + h2 * rhs[idx]) / 4.0;

                    max_diff = max_diff.max((u_new[idx] - u[idx]).abs());
                }
            }

            // Swap buffers
            std::mem::swap(&mut u, &mut u_new);

            iterations = iter + 1;

            // Check convergence
            if max_diff < self.tolerance {
                break;
            }
        }

        (u, iterations)
    }

    /// Solve using Gauss-Seidel iteration (faster convergence than Jacobi).
    #[must_use]
    pub fn solve_gauss_seidel(&self, rhs: &[f64]) -> (Vec<f64>, usize) {
        let n = self.resolution;
        let h2 = self.h * self.h;

        // Initialize solution to zero
        let mut u = vec![0.0; n * n];

        let mut iterations = 0;

        for iter in 0..self.max_iterations {
            let mut max_diff = 0.0f64;

            // Update interior points (in-place)
            for i in 1..n - 1 {
                for j in 1..n - 1 {
                    let idx = i * n + j;
                    let left = u[(i - 1) * n + j];
                    let right = u[(i + 1) * n + j];
                    let down = u[i * n + (j - 1)];
                    let up = u[i * n + (j + 1)];

                    let u_old = u[idx];
                    u[idx] = (left + right + down + up + h2 * rhs[idx]) / 4.0;

                    max_diff = max_diff.max((u[idx] - u_old).abs());
                }
            }

            iterations = iter + 1;

            // Check convergence
            if max_diff < self.tolerance {
                break;
            }
        }

        (u, iterations)
    }

    /// Solve using Successive Over-Relaxation (SOR) - even faster convergence.
    ///
    /// Optimal omega for Poisson on unit square: ω = 2 / (1 + sin(π*h))
    #[must_use]
    pub fn solve_sor(&self, rhs: &[f64], omega: Option<f64>) -> (Vec<f64>, usize) {
        let n = self.resolution;
        let h2 = self.h * self.h;

        // Optimal omega for this problem
        let omega = omega.unwrap_or_else(|| 2.0 / (1.0 + (std::f64::consts::PI * self.h).sin()));

        // Initialize solution to zero
        let mut u = vec![0.0; n * n];

        let mut iterations = 0;

        for iter in 0..self.max_iterations {
            let mut max_diff = 0.0f64;

            // Red-black ordering for better parallelism (though we do serial here)
            for color in 0..2 {
                for i in 1..n - 1 {
                    for j in 1..n - 1 {
                        // Red-black check
                        if (i + j) % 2 != color {
                            continue;
                        }

                        let idx = i * n + j;
                        let left = u[(i - 1) * n + j];
                        let right = u[(i + 1) * n + j];
                        let down = u[i * n + (j - 1)];
                        let up = u[i * n + (j + 1)];

                        let u_old = u[idx];
                        let u_gs = (left + right + down + up + h2 * rhs[idx]) / 4.0;
                        u[idx] = u_old + omega * (u_gs - u_old);

                        max_diff = max_diff.max((u[idx] - u_old).abs());
                    }
                }
            }

            iterations = iter + 1;

            // Check convergence
            if max_diff < self.tolerance {
                break;
            }
        }

        (u, iterations)
    }

    /// Estimate memory usage in MB
    #[must_use]
    pub fn memory_usage_mb(&self) -> f64 {
        // Two N×N arrays of f64 (8 bytes each)
        2.0 * (self.resolution * self.resolution * 8) as f64 / (1024.0 * 1024.0)
    }
}

// ============================================================================
// Finite Element Method (FEM) Solver
// ============================================================================

/// FEM solver using linear (P1) elements for 2D Poisson equation.
///
/// Solves: -∇²u = f on [0,1]² with Dirichlet boundary conditions u = 0.
///
/// Uses bilinear basis functions on a uniform grid, which reduces to a
/// similar stencil as FDM but with slightly different weights.
pub struct FemSolver {
    /// Grid resolution (N x N nodes)
    resolution: usize,
    /// Grid spacing h = 1/(N-1)
    h: f64,
    /// Maximum iterations for CG solver
    max_iterations: usize,
    /// Convergence tolerance
    tolerance: f64,
}

impl FemSolver {
    /// Create a new FEM solver
    #[must_use]
    pub fn new(resolution: usize) -> Self {
        let h = 1.0 / (resolution - 1) as f64;
        Self {
            resolution,
            h,
            max_iterations: 10000,
            tolerance: 1e-6,
        }
    }

    /// Set maximum iterations
    #[must_use]
    pub fn with_max_iterations(mut self, max_iter: usize) -> Self {
        self.max_iterations = max_iter;
        self
    }

    /// Set convergence tolerance
    #[must_use]
    pub fn with_tolerance(mut self, tol: f64) -> Self {
        self.tolerance = tol;
        self
    }

    /// Solve using Conjugate Gradient method.
    ///
    /// For the Poisson equation with P1 elements on a uniform grid,
    /// the stiffness matrix A has a known structure that we can exploit.
    ///
    /// # Arguments
    /// * `rhs` - Right-hand side (mass matrix times f, integrated)
    ///
    /// # Returns
    /// Tuple of (solution, iterations)
    #[must_use]
    pub fn solve_cg(&self, rhs: &[f64]) -> (Vec<f64>, usize) {
        let n = self.resolution;

        // For P1 FEM on uniform grid, the load vector needs mass matrix weighting
        // For simplicity, we use the same discretization as FDM (point evaluation)
        // A proper FEM would integrate f against basis functions

        // Initialize solution
        let mut x = vec![0.0; n * n];

        // Compute initial residual: r = b - A*x (x=0, so r = b)
        let mut r = rhs.to_vec();

        // Apply boundary conditions to residual
        self.apply_boundary_conditions(&mut r);

        // Initial search direction
        let mut p = r.clone();

        // Compute r·r
        let mut r_dot_r = self.dot_product(&r, &r);

        if r_dot_r.sqrt() < self.tolerance {
            return (x, 0);
        }

        let mut iterations = 0;

        for iter in 0..self.max_iterations {
            // Compute A*p
            let ap = self.apply_stiffness(&p);

            // Compute step size: α = (r·r) / (p·Ap)
            let p_dot_ap = self.dot_product(&p, &ap);
            if p_dot_ap.abs() < 1e-15 {
                break;
            }
            let alpha = r_dot_r / p_dot_ap;

            // Update solution: x = x + α*p
            for i in 0..x.len() {
                x[i] += alpha * p[i];
            }

            // Update residual: r = r - α*Ap
            for i in 0..r.len() {
                r[i] -= alpha * ap[i];
            }

            // Apply boundary conditions
            self.apply_boundary_conditions(&mut r);

            // Check convergence
            let r_dot_r_new = self.dot_product(&r, &r);
            iterations = iter + 1;

            if r_dot_r_new.sqrt() < self.tolerance {
                break;
            }

            // Compute β = (r_new·r_new) / (r·r)
            let beta = r_dot_r_new / r_dot_r;
            r_dot_r = r_dot_r_new;

            // Update search direction: p = r + β*p
            for i in 0..p.len() {
                p[i] = r[i] + beta * p[i];
            }
        }

        (x, iterations)
    }

    /// Apply stiffness matrix to a vector (matrix-free)
    ///
    /// For P1 elements on uniform grid, A*u at interior point (i,j) is:
    /// (4*u[i,j] - u[i-1,j] - u[i+1,j] - u[i,j-1] - u[i,j+1]) / h²
    fn apply_stiffness(&self, u: &[f64]) -> Vec<f64> {
        let n = self.resolution;
        let h2 = self.h * self.h;
        let mut result = vec![0.0; n * n];

        for i in 1..n - 1 {
            for j in 1..n - 1 {
                let idx = i * n + j;
                let center = u[idx];
                let left = u[(i - 1) * n + j];
                let right = u[(i + 1) * n + j];
                let down = u[i * n + (j - 1)];
                let up = u[i * n + (j + 1)];

                result[idx] = (4.0 * center - left - right - down - up) / h2;
            }
        }

        result
    }

    /// Apply Dirichlet boundary conditions (zero on boundary)
    fn apply_boundary_conditions(&self, v: &mut [f64]) {
        let n = self.resolution;

        // Top and bottom boundaries
        for j in 0..n {
            v[j] = 0.0; // i = 0
            v[(n - 1) * n + j] = 0.0; // i = n-1
        }

        // Left and right boundaries
        for i in 0..n {
            v[i * n] = 0.0; // j = 0
            v[i * n + (n - 1)] = 0.0; // j = n-1
        }
    }

    /// Compute dot product
    fn dot_product(&self, a: &[f64], b: &[f64]) -> f64 {
        a.iter().zip(b.iter()).map(|(x, y)| x * y).sum()
    }

    /// Estimate memory usage in MB
    #[must_use]
    pub fn memory_usage_mb(&self) -> f64 {
        // Four N×N arrays of f64: x, r, p, ap
        4.0 * (self.resolution * self.resolution * 8) as f64 / (1024.0 * 1024.0)
    }
}

// ============================================================================
// Benchmark Runner
// ============================================================================

/// Benchmark runner that compares FNO, FEM, and FDM solvers
pub struct BenchmarkRunner {
    /// Configuration
    config: BenchmarkConfig,
}

impl BenchmarkRunner {
    /// Create a new benchmark runner
    #[must_use]
    pub fn new(config: BenchmarkConfig) -> Self {
        Self { config }
    }

    /// Run benchmarks for FDM and FEM solvers
    ///
    /// Note: FNO benchmarks require a trained model and are run separately.
    #[must_use]
    pub fn run_classical_benchmarks(&self) -> Vec<BenchmarkResult> {
        let mut results = Vec::new();

        for &resolution in &self.config.resolutions {
            // Generate test problem
            let rhs = self.generate_test_problem(resolution);

            // Benchmark FDM (Gauss-Seidel)
            let fdm_result = self.benchmark_fdm(resolution, &rhs);
            results.push(fdm_result);

            // Benchmark FEM (CG)
            let fem_result = self.benchmark_fem(resolution, &rhs);
            results.push(fem_result);
        }

        results
    }

    /// Benchmark FDM solver
    fn benchmark_fdm(&self, resolution: usize, rhs: &[f64]) -> BenchmarkResult {
        let solver = FdmSolver::new(resolution)
            .with_max_iterations(50000)
            .with_tolerance(1e-8);

        let start = Instant::now();
        let (solution, iterations) = solver.solve_sor(rhs, None);
        let elapsed = start.elapsed();

        let mut result = BenchmarkResult::new("FDM-SOR", resolution, self.config.pde_type)
            .with_time(elapsed.as_secs_f64() * 1000.0)
            .with_memory(solver.memory_usage_mb())
            .with_iterations(iterations);

        // Compute error if we have a reference solution
        if let Some(error) = self.compute_l2_error(&solution, resolution) {
            result = result.with_l2_error(error);
        }

        result
    }

    /// Benchmark FEM solver
    fn benchmark_fem(&self, resolution: usize, rhs: &[f64]) -> BenchmarkResult {
        let solver = FemSolver::new(resolution)
            .with_max_iterations(50000)
            .with_tolerance(1e-8);

        let start = Instant::now();
        let (solution, iterations) = solver.solve_cg(rhs);
        let elapsed = start.elapsed();

        let mut result = BenchmarkResult::new("FEM-CG", resolution, self.config.pde_type)
            .with_time(elapsed.as_secs_f64() * 1000.0)
            .with_memory(solver.memory_usage_mb())
            .with_iterations(iterations);

        // Compute error if we have a reference solution
        if let Some(error) = self.compute_l2_error(&solution, resolution) {
            result = result.with_l2_error(error);
        }

        result
    }

    /// Generate a test problem (right-hand side f)
    ///
    /// For Poisson equation, we use f(x,y) = 2π²sin(πx)sin(πy)
    /// which has exact solution u(x,y) = sin(πx)sin(πy)
    fn generate_test_problem(&self, resolution: usize) -> Vec<f64> {
        let n = resolution;
        let h = 1.0 / (n - 1) as f64;
        let mut rhs = vec![0.0; n * n];

        let pi = std::f64::consts::PI;

        for i in 0..n {
            for j in 0..n {
                let x = i as f64 * h;
                let y = j as f64 * h;

                // f = 2π²sin(πx)sin(πy) for exact solution u = sin(πx)sin(πy)
                rhs[i * n + j] = 2.0 * pi * pi * (pi * x).sin() * (pi * y).sin();
            }
        }

        rhs
    }

    /// Compute L2 error vs analytical solution
    ///
    /// Exact solution: u(x,y) = sin(πx)sin(πy)
    fn compute_l2_error(&self, solution: &[f64], resolution: usize) -> Option<f64> {
        let n = resolution;
        let h = 1.0 / (n - 1) as f64;
        let pi = std::f64::consts::PI;

        let mut error_sum = 0.0;

        for i in 0..n {
            for j in 0..n {
                let x = i as f64 * h;
                let y = j as f64 * h;

                let exact = (pi * x).sin() * (pi * y).sin();
                let numerical = solution[i * n + j];

                error_sum += (exact - numerical).powi(2);
            }
        }

        Some((error_sum * h * h).sqrt()) // L2 norm with grid weighting
    }

    /// Run multiple benchmark iterations and compute summary statistics
    #[must_use]
    pub fn run_with_statistics(&self, n_runs: usize) -> Vec<BenchmarkSummary> {
        let mut all_results: Vec<Vec<BenchmarkResult>> = Vec::new();

        for _ in 0..n_runs {
            let results = self.run_classical_benchmarks();
            all_results.push(results);
        }

        // Group results by (method, resolution)
        let mut grouped: std::collections::HashMap<(String, usize), Vec<BenchmarkResult>> =
            std::collections::HashMap::new();

        for run_results in all_results {
            for result in run_results {
                let key = (result.method.clone(), result.resolution);
                grouped.entry(key).or_default().push(result);
            }
        }

        // Compute summaries
        grouped
            .values()
            .filter_map(|results| BenchmarkSummary::from_results(results))
            .collect()
    }

    /// Print benchmark results in a formatted table
    pub fn print_results(results: &[BenchmarkResult]) {
        println!("\n{:=<80}", "");
        println!("PDE Solver Benchmark Results");
        println!("{:=<80}\n", "");

        println!(
            "{:<12} {:>10} {:>12} {:>12} {:>12} {:>10}",
            "Method", "Resolution", "Time (ms)", "L2 Error", "Memory (MB)", "Iterations"
        );
        println!("{:-<80}", "");

        for result in results {
            let l2_str = result
                .l2_error
                .map_or_else(|| "N/A".to_string(), |e| format!("{e:.2e}"));
            let iter_str = result
                .iterations
                .map_or_else(|| "N/A".to_string(), |i| i.to_string());

            println!(
                "{:<12} {:>10} {:>12.2} {:>12} {:>12.2} {:>10}",
                result.method,
                result.resolution,
                result.solve_time_ms,
                l2_str,
                result.memory_mb,
                iter_str
            );
        }

        println!();
    }

    /// Print summary statistics
    pub fn print_summary(summaries: &[BenchmarkSummary]) {
        println!("\n{:=<80}", "");
        println!("Benchmark Summary Statistics");
        println!("{:=<80}\n", "");

        println!(
            "{:<12} {:>10} {:>12} {:>12} {:>12}",
            "Method", "Resolution", "Avg Time", "Std Dev", "Avg L2 Error"
        );
        println!("{:-<80}", "");

        for summary in summaries {
            let l2_str = summary
                .avg_l2_error
                .map_or_else(|| "N/A".to_string(), |e| format!("{e:.2e}"));

            println!(
                "{:<12} {:>10} {:>10.2}ms {:>10.2}ms {:>12}",
                summary.method,
                summary.resolution,
                summary.avg_time_ms,
                summary.std_time_ms,
                l2_str
            );
        }

        println!();
    }

    /// Benchmark FNO inference
    ///
    /// Takes a trained FNO model and benchmarks its inference time.
    /// The model should be configured for the same PDE type as the benchmark.
    ///
    /// # Arguments
    /// * `model` - A trained `FNO2d` model
    /// * `resolution` - Grid resolution to benchmark
    ///
    /// # Returns
    /// Benchmark result with timing and error metrics
    #[must_use]
    pub fn benchmark_fno(&self, model: &FNO2d<CpuBackend>, resolution: usize) -> BenchmarkResult {
        let device = CpuDevice::default();

        // Generate input tensor (coefficient field for the PDE)
        // Shape: [1, 1, resolution, resolution]
        let input_data: Vec<f32> = (0..resolution * resolution)
            .map(|i| {
                let x = (i % resolution) as f32 / (resolution - 1) as f32;
                let y = (i / resolution) as f32 / (resolution - 1) as f32;
                // Use sinusoidal input for testing
                (std::f32::consts::PI * x).sin() * (std::f32::consts::PI * y).sin()
            })
            .collect();

        let input = GenericTensor::<CpuBackend, 4>::from_slice(
            &input_data,
            [1, 1, resolution, resolution],
            &device,
        );

        // Warm-up run
        let _ = model.forward_4d(&input);

        // Timed runs - multiple iterations for more accurate timing
        let n_iterations = 10;
        let start = Instant::now();
        let mut output = None;
        for _ in 0..n_iterations {
            output = Some(model.forward_4d(&input));
        }
        let elapsed = start.elapsed();
        let avg_time_ms = elapsed.as_secs_f64() * 1000.0 / f64::from(n_iterations);

        // Extract output for error computation
        let output = output.unwrap();
        let output_data = output.to_vec();

        // Compute L2 error vs analytical solution (same as classical solvers)
        let l2_error = self.compute_fno_l2_error(&output_data, resolution);

        // Estimate memory: input + output + model parameters (rough estimate)
        // Input: 1 x 1 x res x res x 4 bytes
        // Output: same
        // Model: typically ~1-10MB for small FNO
        let io_memory = 2.0 * (resolution * resolution * 4) as f64 / (1024.0 * 1024.0);
        let model_memory_estimate = 5.0; // MB, rough estimate
        let total_memory = io_memory + model_memory_estimate;

        let mut result = BenchmarkResult::new("FNO", resolution, self.config.pde_type)
            .with_time(avg_time_ms)
            .with_memory(total_memory);

        if let Some(error) = l2_error {
            result = result.with_l2_error(error);
        }

        result
    }

    /// Compute L2 error for FNO output vs analytical solution
    fn compute_fno_l2_error(&self, output: &[f32], resolution: usize) -> Option<f64> {
        // Output shape is [1, 1, resolution, resolution], flattened
        // We expect output.len() == resolution * resolution
        if output.len() < resolution * resolution {
            return None;
        }

        let n = resolution;
        let h = 1.0 / (n - 1) as f64;
        let pi = std::f64::consts::PI;

        let mut error_sum = 0.0;

        for i in 0..n {
            for j in 0..n {
                let x = i as f64 * h;
                let y = j as f64 * h;

                // Exact solution for Poisson with f = 2π²sin(πx)sin(πy)
                let exact = (pi * x).sin() * (pi * y).sin();
                let numerical = f64::from(output[i * n + j]);

                error_sum += (exact - numerical).powi(2);
            }
        }

        Some((error_sum * h * h).sqrt())
    }

    /// Run all benchmarks including FNO
    ///
    /// This runs both classical (FDM/FEM) and neural operator (FNO) benchmarks
    /// for comparison.
    ///
    /// # Arguments
    /// * `model` - Optional trained FNO model. If None, only classical benchmarks are run.
    ///
    /// # Returns
    /// Vector of benchmark results for all methods and resolutions
    #[must_use]
    pub fn run_all_benchmarks(&self, model: Option<&FNO2d<CpuBackend>>) -> Vec<BenchmarkResult> {
        let mut results = Vec::new();

        for &resolution in &self.config.resolutions {
            // Generate test problem
            let rhs = self.generate_test_problem(resolution);

            // Benchmark classical methods
            let fdm_result = self.benchmark_fdm(resolution, &rhs);
            results.push(fdm_result);

            let fem_result = self.benchmark_fem(resolution, &rhs);
            results.push(fem_result);

            // Benchmark FNO if model is provided
            if let Some(fno_model) = model {
                let fno_result = self.benchmark_fno(fno_model, resolution);
                results.push(fno_result);
            }
        }

        results
    }

    /// Compute speedup factors comparing FNO to classical methods
    ///
    /// Returns a formatted string summarizing speedup results.
    #[must_use]
    pub fn compute_speedup_report(results: &[BenchmarkResult]) -> String {
        let mut report = String::new();
        report.push('\n');
        report.push_str(&"=".repeat(80));
        report.push_str("\nFNO Speedup Analysis\n");
        report.push_str(&"=".repeat(80));
        report.push_str("\n\n");

        // Group results by resolution
        let mut by_resolution: std::collections::HashMap<usize, Vec<&BenchmarkResult>> =
            std::collections::HashMap::new();

        for result in results {
            by_resolution
                .entry(result.resolution)
                .or_default()
                .push(result);
        }

        for (resolution, res_results) in &by_resolution {
            let fno_time = res_results
                .iter()
                .find(|r| r.method == "FNO")
                .map(|r| r.solve_time_ms);

            if let Some(fno_t) = fno_time {
                report.push_str(&format!("Resolution: {resolution}x{resolution}\n"));
                report.push_str(&"-".repeat(40));
                report.push('\n');

                for result in res_results {
                    if result.method != "FNO" {
                        let speedup = result.solve_time_ms / fno_t;
                        report.push_str(&format!(
                            "  FNO vs {}: {:.1}x speedup ({:.2}ms vs {:.2}ms)\n",
                            result.method, speedup, fno_t, result.solve_time_ms
                        ));
                    }
                }
                report.push('\n');
            }
        }

        report
    }
}

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

    #[test]
    fn test_fdm_solver_convergence() {
        let resolution = 32;
        let solver = FdmSolver::new(resolution).with_tolerance(1e-6);

        // Generate test problem with known solution
        let n = resolution;
        let h = 1.0 / (n - 1) as f64;
        let pi = std::f64::consts::PI;

        let mut rhs = vec![0.0; n * n];
        for i in 0..n {
            for j in 0..n {
                let x = i as f64 * h;
                let y = j as f64 * h;
                rhs[i * n + j] = 2.0 * pi * pi * (pi * x).sin() * (pi * y).sin();
            }
        }

        let (solution, iterations) = solver.solve_sor(&rhs, None);

        // Check that it converged
        assert!(iterations < 10000, "Should converge within max iterations");

        // Check error at center point
        let center = n / 2;
        let x = center as f64 * h;
        let y = center as f64 * h;
        let exact = (pi * x).sin() * (pi * y).sin();
        let numerical = solution[center * n + center];

        assert!(
            (exact - numerical).abs() < 0.01,
            "Solution should be close to exact at center: exact={}, numerical={}",
            exact,
            numerical
        );
    }

    #[test]
    fn test_fem_solver_convergence() {
        let resolution = 32;
        let solver = FemSolver::new(resolution).with_tolerance(1e-6);

        // Generate test problem with known solution
        let n = resolution;
        let h = 1.0 / (n - 1) as f64;
        let pi = std::f64::consts::PI;

        let mut rhs = vec![0.0; n * n];
        for i in 0..n {
            for j in 0..n {
                let x = i as f64 * h;
                let y = j as f64 * h;
                rhs[i * n + j] = 2.0 * pi * pi * (pi * x).sin() * (pi * y).sin();
            }
        }

        let (solution, iterations) = solver.solve_cg(&rhs);

        // Check that it converged
        assert!(iterations < 10000, "Should converge within max iterations");

        // Check error at center point
        let center = n / 2;
        let x = center as f64 * h;
        let y = center as f64 * h;
        let exact = (pi * x).sin() * (pi * y).sin();
        let numerical = solution[center * n + center];

        assert!(
            (exact - numerical).abs() < 0.01,
            "Solution should be close to exact at center: exact={}, numerical={}",
            exact,
            numerical
        );
    }

    #[test]
    fn test_benchmark_runner() {
        let config = BenchmarkConfig::quick();
        let runner = BenchmarkRunner::new(config);

        let results = runner.run_classical_benchmarks();

        // Should have results for each resolution and method
        assert!(!results.is_empty(), "Should have benchmark results");

        // Check that all results have valid times
        for result in &results {
            assert!(result.solve_time_ms > 0.0, "Solve time should be positive");
            assert!(result.memory_mb > 0.0, "Memory usage should be positive");
        }
    }

    #[test]
    fn test_benchmark_summary() {
        let results = vec![
            BenchmarkResult::new("FDM", 32, PDEType::Poisson)
                .with_time(10.0)
                .with_l2_error(0.01),
            BenchmarkResult::new("FDM", 32, PDEType::Poisson)
                .with_time(12.0)
                .with_l2_error(0.01),
            BenchmarkResult::new("FDM", 32, PDEType::Poisson)
                .with_time(11.0)
                .with_l2_error(0.01),
        ];

        let summary = BenchmarkSummary::from_results(&results).unwrap();

        assert_eq!(summary.method, "FDM");
        assert_eq!(summary.resolution, 32);
        assert!((summary.avg_time_ms - 11.0).abs() < 0.01);
        assert_eq!(summary.n_runs, 3);
    }

    #[test]
    fn test_fdm_methods_comparison() {
        // Compare Jacobi, Gauss-Seidel, and SOR convergence rates
        let resolution = 32;
        let n = resolution;
        let h = 1.0 / (n - 1) as f64;
        let pi = std::f64::consts::PI;

        let mut rhs = vec![0.0; n * n];
        for i in 0..n {
            for j in 0..n {
                let x = i as f64 * h;
                let y = j as f64 * h;
                rhs[i * n + j] = 2.0 * pi * pi * (pi * x).sin() * (pi * y).sin();
            }
        }

        let solver = FdmSolver::new(resolution).with_tolerance(1e-6);

        let (_, jacobi_iters) = solver.solve_jacobi(&rhs);
        let (_, gs_iters) = solver.solve_gauss_seidel(&rhs);
        let (_, sor_iters) = solver.solve_sor(&rhs, None);

        // SOR should converge faster than Gauss-Seidel, which should be faster than Jacobi
        assert!(
            sor_iters < gs_iters,
            "SOR ({}) should converge faster than GS ({})",
            sor_iters,
            gs_iters
        );
        assert!(
            gs_iters < jacobi_iters,
            "GS ({}) should converge faster than Jacobi ({})",
            gs_iters,
            jacobi_iters
        );
    }

    #[test]
    fn test_fno_benchmark() {
        // Create a small FNO model for testing
        let device = CpuDevice::default();
        let resolution = 32;

        // FNO with minimal configuration for testing
        // in_channels=1, out_channels=1, width=16, n_modes=8
        let model =
            FNO2d::<CpuBackend>::new(1, 1, 16, 8, &device).expect("Failed to create FNO model");

        // Create benchmark runner
        let config = BenchmarkConfig {
            resolutions: vec![resolution],
            n_problems: 1,
            pde_type: PDEType::Poisson,
            use_reference: false,
            reference_resolution: 64,
        };
        let runner = BenchmarkRunner::new(config);

        // Run FNO benchmark
        let result = runner.benchmark_fno(&model, resolution);

        // Verify result
        assert_eq!(result.method, "FNO");
        assert_eq!(result.resolution, resolution);
        assert!(
            result.solve_time_ms > 0.0,
            "FNO should have positive solve time"
        );
        assert!(
            result.memory_mb > 0.0,
            "FNO should have positive memory usage"
        );
        // L2 error will be large for untrained model, but should exist
        assert!(
            result.l2_error.is_some(),
            "FNO benchmark should compute L2 error"
        );
    }

    #[test]
    fn test_run_all_benchmarks_with_fno() {
        let device = CpuDevice::default();

        // Create small FNO model
        let model =
            FNO2d::<CpuBackend>::new(1, 1, 16, 8, &device).expect("Failed to create FNO model");

        // Quick config for testing
        let config = BenchmarkConfig {
            resolutions: vec![16], // Small resolution for fast test
            n_problems: 1,
            pde_type: PDEType::Poisson,
            use_reference: false,
            reference_resolution: 32,
        };
        let runner = BenchmarkRunner::new(config);

        // Run all benchmarks including FNO
        let results = runner.run_all_benchmarks(Some(&model));

        // Should have 3 results: FDM, FEM, FNO
        assert_eq!(results.len(), 3, "Should have FDM, FEM, and FNO results");

        // Verify we have one of each method
        let methods: Vec<&str> = results.iter().map(|r| r.method.as_str()).collect();
        assert!(methods.contains(&"FDM-SOR"), "Should have FDM result");
        assert!(methods.contains(&"FEM-CG"), "Should have FEM result");
        assert!(methods.contains(&"FNO"), "Should have FNO result");
    }

    #[test]
    fn test_speedup_report() {
        // Create mock benchmark results
        let results = vec![
            BenchmarkResult::new("FNO", 64, PDEType::Poisson).with_time(1.0),
            BenchmarkResult::new("FDM-SOR", 64, PDEType::Poisson).with_time(100.0),
            BenchmarkResult::new("FEM-CG", 64, PDEType::Poisson).with_time(50.0),
        ];

        let report = BenchmarkRunner::compute_speedup_report(&results);

        // Report should contain speedup information
        assert!(
            report.contains("FNO vs FDM-SOR"),
            "Report should compare FNO vs FDM"
        );
        assert!(
            report.contains("FNO vs FEM-CG"),
            "Report should compare FNO vs FEM"
        );
        assert!(
            report.contains("100.0x"),
            "Should show ~100x speedup vs FDM"
        );
        assert!(report.contains("50.0x"), "Should show ~50x speedup vs FEM");
    }
}