//! Classical PDE solvers for comparison.
//!
//! Provides finite difference solvers to compare against neural operators.

use neuralop_studio_shared::{BCType, PDEDefinition, PDEType, SolutionField};

/// Classical PDE solver.
#[derive(Debug)]
pub struct ClassicalSolver {
    /// Maximum iterations.
    max_iterations: usize,
    /// Convergence tolerance.
    tolerance: f64,
}

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

impl ClassicalSolver {
    /// Create a new classical solver.
    #[must_use]
    pub fn new() -> Self {
        Self {
            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 the PDE.
    pub fn solve(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
        match pde.pde_type {
            PDEType::Poisson => self.solve_poisson(pde),
            PDEType::Heat => self.solve_heat(pde),
            PDEType::Wave => self.solve_wave(pde),
            PDEType::Burgers => self.solve_burgers(pde),
            _ => Err(format!("Unsupported PDE type: {:?}", pde.pde_type)),
        }
    }

    /// Solve Poisson equation using Gauss-Seidel iteration.
    fn solve_poisson(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
        let nx = pde.domain.resolution[0];
        let ny = if pde.domain.dimensions > 1 {
            pde.domain.resolution.get(1).copied().unwrap_or(1)
        } else {
            1
        };

        let x_range = pde.domain.bounds[0];
        let y_range = if pde.domain.dimensions > 1 {
            pde.domain.bounds.get(1).copied().unwrap_or((0.0, 1.0))
        } else {
            (0.0, 1.0)
        };

        let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
        let dy = if ny > 1 {
            (y_range.1 - y_range.0) / (ny - 1) as f64
        } else {
            dx
        };

        let mut u = vec![0.0; nx * ny];
        let mut u_new = vec![0.0; nx * ny];

        // Apply boundary conditions
        self.apply_boundary_conditions(&mut u, pde, nx, ny);

        // Source term
        let source: Vec<f64> = (0..nx * ny)
            .map(|idx| {
                let i = idx % nx;
                let j = idx / nx;
                let x = x_range.0 + i as f64 * dx;
                let y = y_range.0 + j as f64 * dy;
                self.evaluate_source(pde, x, y)
            })
            .collect();

        // Gauss-Seidel iteration
        for iteration in 0..self.max_iterations {
            u_new.copy_from_slice(&u);

            for j in 1..ny.saturating_sub(1).max(1) {
                for i in 1..nx.saturating_sub(1) {
                    let idx = j * nx + i;
                    let idx_left = j * nx + i - 1;
                    let idx_right = j * nx + i + 1;
                    let idx_down = (j.saturating_sub(1)) * nx + i;
                    let idx_up = (j + 1).min(ny - 1) * nx + i;

                    if ny == 1 {
                        // 1D case
                        u_new[idx] = 0.5 * (u[idx_left] + u[idx_right] + dx * dx * source[idx]);
                    } else {
                        // 2D case
                        u_new[idx] = 0.25
                            * (u[idx_left]
                                + u_new[idx_right.min(nx * ny - 1)]
                                + u[idx_down]
                                + u_new[idx_up.min(nx * ny - 1)]
                                + dx * dy * source[idx]);
                    }
                }
            }

            // Check convergence
            let diff: f64 = u
                .iter()
                .zip(u_new.iter())
                .map(|(a, b)| (a - b).abs())
                .sum::<f64>()
                / (nx * ny) as f64;

            u.copy_from_slice(&u_new);

            if diff < self.tolerance {
                break;
            }

            if iteration == self.max_iterations - 1 {
                // Warning: did not converge, but continue
            }
        }

        // Generate coordinates
        let coordinates: Vec<Vec<f64>> = (0..nx * ny)
            .map(|idx| {
                let i = idx % nx;
                let j = idx / nx;
                let x = x_range.0 + i as f64 * dx;
                let y = y_range.0 + j as f64 * dy;
                vec![x, y]
            })
            .collect();

        Ok(SolutionField {
            coordinates,
            values: u,
            name: "u".to_string(),
            time: None,
        })
    }

    /// Solve heat equation using explicit finite differences.
    fn solve_heat(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
        let nx = pde.domain.resolution[0];
        let (t_start, t_end) = pde.domain.time_bounds.unwrap_or((0.0, 1.0));
        let nt = pde.domain.time_resolution.unwrap_or(100);

        let x_range = pde.domain.bounds[0];
        let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
        let dt = (t_end - t_start) / nt as f64;

        let alpha = pde.parameters.diffusion.unwrap_or(0.01);

        // CFL condition check
        let cfl = alpha * dt / (dx * dx);
        if cfl > 0.5 {
            // Adjust dt if unstable
            let _dt = 0.4 * dx * dx / alpha;
        }

        let mut u = vec![0.0; nx];
        let mut u_new = vec![0.0; nx];

        // Apply initial condition
        if let Some(ref ic) = pde.initial_condition {
            for i in 0..nx {
                let x = x_range.0 + i as f64 * dx;
                u[i] = self.evaluate_initial(ic, x);
            }
        }

        // Time stepping
        for _ in 0..nt {
            u_new.copy_from_slice(&u);

            for i in 1..nx - 1 {
                let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
                u_new[i] = u[i] + dt * alpha * laplacian;
            }

            // Apply boundary conditions
            u_new[0] = 0.0;
            u_new[nx - 1] = 0.0;

            u.copy_from_slice(&u_new);
        }

        let coordinates: Vec<Vec<f64>> = (0..nx)
            .map(|i| {
                let x = x_range.0 + i as f64 * dx;
                vec![x]
            })
            .collect();

        Ok(SolutionField {
            coordinates,
            values: u,
            name: "u".to_string(),
            time: Some(t_end),
        })
    }

    /// Solve wave equation using explicit finite differences.
    fn solve_wave(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
        let nx = pde.domain.resolution[0];
        let (t_start, t_end) = pde.domain.time_bounds.unwrap_or((0.0, 1.0));
        let nt = pde.domain.time_resolution.unwrap_or(100);

        let x_range = pde.domain.bounds[0];
        let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
        let dt = (t_end - t_start) / nt as f64;

        let c = pde.parameters.wave_speed.unwrap_or(1.0);
        let courant = c * dt / dx;

        if courant > 1.0 {
            // Warn about stability
        }

        let mut u = vec![0.0; nx];
        let mut u_prev = vec![0.0; nx];
        let mut u_new = vec![0.0; nx];

        // Initial conditions
        if let Some(ref ic) = pde.initial_condition {
            for i in 0..nx {
                let x = x_range.0 + i as f64 * dx;
                u[i] = self.evaluate_initial(ic, x);
                u_prev[i] = u[i]; // Zero initial velocity assumed
            }
        }

        // Time stepping
        for _ in 0..nt {
            for i in 1..nx - 1 {
                let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
                u_new[i] = 2.0 * u[i] - u_prev[i] + dt * dt * c * c * laplacian;
            }

            u_new[0] = 0.0;
            u_new[nx - 1] = 0.0;

            u_prev.copy_from_slice(&u);
            u.copy_from_slice(&u_new);
        }

        let coordinates: Vec<Vec<f64>> = (0..nx).map(|i| vec![x_range.0 + i as f64 * dx]).collect();

        Ok(SolutionField {
            coordinates,
            values: u,
            name: "u".to_string(),
            time: Some(t_end),
        })
    }

    /// Solve Burgers' equation using Lax-Wendroff scheme.
    fn solve_burgers(&self, pde: &PDEDefinition) -> Result<SolutionField, String> {
        let nx = pde.domain.resolution[0];
        let (t_start, t_end) = pde.domain.time_bounds.unwrap_or((0.0, 1.0));
        let nt = pde.domain.time_resolution.unwrap_or(100);

        let x_range = pde.domain.bounds[0];
        let dx = (x_range.1 - x_range.0) / (nx - 1) as f64;
        let dt = (t_end - t_start) / nt as f64;

        let nu = pde.parameters.diffusion.unwrap_or(0.01);

        let mut u = vec![0.0; nx];
        let mut u_new = vec![0.0; nx];

        // Initial condition
        if let Some(ref ic) = pde.initial_condition {
            for i in 0..nx {
                let x = x_range.0 + i as f64 * dx;
                u[i] = self.evaluate_initial(ic, x);
            }
        }

        // Time stepping
        for _ in 0..nt {
            u_new.copy_from_slice(&u);

            for i in 1..nx - 1 {
                let laplacian = (u[i + 1] - 2.0 * u[i] + u[i - 1]) / (dx * dx);
                let du_dx = (u[i + 1] - u[i - 1]) / (2.0 * dx);
                let convection = u[i] * du_dx;

                u_new[i] = u[i] + dt * (nu * laplacian - convection);
            }

            // Periodic boundary conditions
            u_new[0] = u_new[nx - 2];
            u_new[nx - 1] = u_new[1];

            u.copy_from_slice(&u_new);
        }

        let coordinates: Vec<Vec<f64>> = (0..nx).map(|i| vec![x_range.0 + i as f64 * dx]).collect();

        Ok(SolutionField {
            coordinates,
            values: u,
            name: "u".to_string(),
            time: Some(t_end),
        })
    }

    /// Apply boundary conditions.
    fn apply_boundary_conditions(&self, u: &mut [f64], pde: &PDEDefinition, nx: usize, ny: usize) {
        for bc in &pde.boundary_conditions {
            match bc.bc_type {
                BCType::Dirichlet => {
                    let value: f64 = bc.value.parse().unwrap_or(0.0);

                    match bc.boundary.as_str() {
                        "left" => {
                            for j in 0..ny {
                                u[j * nx] = value;
                            }
                        }
                        "right" => {
                            for j in 0..ny {
                                u[j * nx + nx - 1] = value;
                            }
                        }
                        "bottom" => {
                            for i in 0..nx {
                                u[i] = value;
                            }
                        }
                        "top" => {
                            for i in 0..nx {
                                u[(ny - 1) * nx + i] = value;
                            }
                        }
                        "all" | _ => {
                            // All boundaries
                            for i in 0..nx {
                                u[i] = value;
                                u[(ny - 1) * nx + i] = value;
                            }
                            for j in 0..ny {
                                u[j * nx] = value;
                                u[j * nx + nx - 1] = value;
                            }
                        }
                    }
                }
                BCType::Periodic => {
                    // Handled in time stepping
                }
                _ => {}
            }
        }
    }

    /// Evaluate source term.
    fn evaluate_source(&self, pde: &PDEDefinition, x: f64, y: f64) -> f64 {
        if let Some(ref source) = pde.parameters.source_term {
            // Simple pattern matching for common sources
            if source.contains("sin") {
                (std::f64::consts::PI * x).sin() * (std::f64::consts::PI * y).sin()
            } else {
                source.parse().unwrap_or(0.0)
            }
        } else {
            0.0
        }
    }

    /// Evaluate initial condition.
    fn evaluate_initial(&self, ic: &neuralop_studio_shared::InitialCondition, x: f64) -> f64 {
        if ic.function.contains("sin") {
            (std::f64::consts::PI * x).sin()
        } else if ic.function.contains("gauss") {
            (-((x - 0.5) * (x - 0.5)) / 0.01).exp()
        } else {
            ic.function.parse().unwrap_or(0.0)
        }
    }
}

/// Compare neural operator solution with classical solution.
#[must_use]
pub fn compare_solutions(neural_solution: &[f64], classical_solution: &[f64]) -> (f64, f64, f64) {
    if neural_solution.is_empty() || classical_solution.is_empty() {
        return (0.0, 0.0, 0.0);
    }

    let min_len = neural_solution.len().min(classical_solution.len());

    let mse: f64 = neural_solution[..min_len]
        .iter()
        .zip(classical_solution[..min_len].iter())
        .map(|(n, c)| (n - c).powi(2))
        .sum::<f64>()
        / min_len as f64;

    let classical_norm: f64 = classical_solution[..min_len]
        .iter()
        .map(|c| c.powi(2))
        .sum::<f64>()
        .sqrt();

    let error_norm: f64 = neural_solution[..min_len]
        .iter()
        .zip(classical_solution[..min_len].iter())
        .map(|(n, c)| (n - c).powi(2))
        .sum::<f64>()
        .sqrt();

    let relative_l2 = if classical_norm > 1e-10 {
        error_norm / classical_norm
    } else {
        error_norm
    };

    let max_error: f64 = neural_solution[..min_len]
        .iter()
        .zip(classical_solution[..min_len].iter())
        .map(|(n, c)| (n - c).abs())
        .fold(0.0, f64::max);

    (mse, relative_l2, max_error)
}

#[cfg(test)]
mod tests {
    use super::*;
    use neuralop_studio_shared::{
        sample_burgers_problem, sample_heat_problem, sample_poisson_problem,
    };

    #[test]
    fn test_solver_creation() {
        let solver = ClassicalSolver::new();
        assert_eq!(solver.max_iterations, 10000);
    }

    #[test]
    fn test_solve_poisson() {
        let solver = ClassicalSolver::new().with_max_iterations(100);
        let pde = sample_poisson_problem();

        let result = solver.solve(&pde);
        assert!(result.is_ok());

        let solution = result.unwrap();
        assert!(!solution.values.is_empty());
    }

    #[test]
    fn test_solve_heat() {
        let solver = ClassicalSolver::new();
        let pde = sample_heat_problem();

        let result = solver.solve(&pde);
        assert!(result.is_ok());

        let solution = result.unwrap();
        assert!(solution.time.is_some());
    }

    #[test]
    fn test_solve_burgers() {
        let solver = ClassicalSolver::new();
        let pde = sample_burgers_problem();

        let result = solver.solve(&pde);
        assert!(result.is_ok());
    }

    #[test]
    fn test_compare_solutions() {
        let neural = vec![0.0, 0.1, 0.2, 0.1, 0.0];
        let classical = vec![0.0, 0.1, 0.2, 0.1, 0.0];

        let (mse, rel_l2, max_err) = compare_solutions(&neural, &classical);
        assert!(mse < 1e-10);
        assert!(rel_l2 < 1e-10);
        assert!(max_err < 1e-10);
    }

    #[test]
    fn test_compare_solutions_with_error() {
        let neural = vec![0.0, 0.11, 0.2, 0.1, 0.0]; // Small error
        let classical = vec![0.0, 0.1, 0.2, 0.1, 0.0];

        let (mse, _rel_l2, max_err) = compare_solutions(&neural, &classical);
        assert!(mse > 0.0);
        assert!((max_err - 0.01).abs() < 1e-10);
    }
}