//! SIMPLE (Semi-Implicit Method for Pressure Linked Equations) algorithm
//!
//! The SIMPLE algorithm is a widely-used iterative solution method for
//! the incompressible Navier-Stokes equations. It uses a pressure-velocity
//! coupling approach to handle the incompressibility constraint.
//!
//! Algorithm steps:
//! 1. Solve momentum equations with guessed pressure field → u*, v*
//! 2. Solve pressure correction equation → p'
//! 3. Correct velocities and pressure
//! 4. Check convergence and iterate

use super::{BoundaryConditions, FlowField, IncompressibleSolver, SolverResult};
use crate::turbulence::{KEpsilonModel, KEpsilonVariant, TurbulenceModel, TurbulenceState};
use crate::{CfdConfig, CfdError, CfdResult};
use async_trait::async_trait;
use nalgebra::{DMatrix, DVector, Vector3};
use std::time::Instant;

/// Parameters for SIMPLE algorithm
#[derive(Debug, Clone)]
pub struct SimpleParameters {
    /// Under-relaxation factor for pressure (typically 0.2-0.8)
    pub pressure_relaxation: f64,
    /// Under-relaxation factor for velocity (typically 0.5-0.8)
    pub velocity_relaxation: f64,
    /// Maximum number of iterations
    pub max_iterations: usize,
    /// Convergence tolerance for residuals
    pub tolerance: f64,
    /// Time step for transient problems
    pub time_step: f64,
    /// Maximum Courant number for stability
    pub max_courant: f64,
    /// Enable turbulence modeling
    pub use_turbulence: bool,
}

impl SimpleParameters {
    /// Create new SIMPLE parameters with default values
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Set pressure under-relaxation factor
    #[must_use]
    pub fn with_pressure_relaxation(mut self, factor: f64) -> Self {
        self.pressure_relaxation = factor.max(0.0); // Ensure non-negative
        self
    }

    /// Set velocity under-relaxation factor
    #[must_use]
    pub fn with_velocity_relaxation(mut self, factor: f64) -> Self {
        self.velocity_relaxation = factor.max(0.0);
        self
    }

    /// 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.abs();
        self
    }

    /// Set time step
    #[must_use]
    pub fn with_time_step(mut self, dt: f64) -> Self {
        self.time_step = dt.abs();
        self
    }

    /// Validate parameters
    pub fn validate(&self) -> CfdResult<()> {
        if self.pressure_relaxation <= 0.0 || self.pressure_relaxation > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Pressure relaxation factor must be in (0, 1]",
            ));
        }

        if self.velocity_relaxation <= 0.0 || self.velocity_relaxation > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Velocity relaxation factor must be in (0, 1]",
            ));
        }

        if self.tolerance <= 0.0 {
            return Err(CfdError::invalid_parameter("Tolerance must be positive"));
        }

        if self.time_step <= 0.0 {
            return Err(CfdError::invalid_parameter("Time step must be positive"));
        }

        Ok(())
    }
}

impl Default for SimpleParameters {
    fn default() -> Self {
        Self {
            pressure_relaxation: 0.3,
            velocity_relaxation: 0.7,
            max_iterations: 1000,
            tolerance: 1e-6,
            time_step: 0.001,
            max_courant: 1.0,
            use_turbulence: false,
        }
    }
}

/// Result of SIMPLE algorithm execution
#[derive(Debug, Clone)]
pub struct SimpleResult {
    /// Base solver result information
    pub solver_result: SolverResult,
    /// Pressure correction iterations per SIMPLE iteration
    pub pressure_iterations: Vec<usize>,
    /// Final mass residual
    pub mass_residual: f64,
    /// Final momentum residual
    pub momentum_residual: f64,
}

/// SIMPLE algorithm implementation
pub struct SimpleSolver {
    /// CFD configuration
    config: CfdConfig,
    /// SIMPLE parameters
    parameters: SimpleParameters,
    /// Linear algebra workspace
    workspace: LinearAlgebraWorkspace,
    /// Turbulence model (optional)
    turbulence_model: Option<KEpsilonModel>,
}

/// Workspace for linear algebra operations
struct LinearAlgebraWorkspace {
    /// Matrix for pressure correction equation
    pressure_matrix: Option<DMatrix<f64>>,
    /// RHS vector for pressure correction
    pressure_rhs: Option<DVector<f64>>,
    /// Solution vector for pressure correction
    pressure_solution: Option<DVector<f64>>,
    /// Momentum equation coefficients
    momentum_coefficients: Option<MomentumCoefficients>,
}

/// Coefficients for momentum equations discretization
#[derive(Debug, Clone)]
struct MomentumCoefficients {
    /// Central coefficient (diagonal)
    pub ap: DMatrix<f64>,
    /// East neighbor coefficient
    pub ae: DMatrix<f64>,
    /// West neighbor coefficient
    pub aw: DMatrix<f64>,
    /// North neighbor coefficient
    pub an: DMatrix<f64>,
    /// South neighbor coefficient
    pub as_: DMatrix<f64>,
    /// Source term
    pub su: DMatrix<f64>,
}

impl SimpleSolver {
    /// Create new SIMPLE solver
    pub fn new(config: CfdConfig, parameters: SimpleParameters) -> CfdResult<Self> {
        config.validate()?;
        parameters.validate()?;

        // Initialize turbulence model if enabled (will be properly sized when flow field is available)
        let turbulence_model = if parameters.use_turbulence {
            None // Will be initialized when flow field dimensions are known
        } else {
            None
        };

        Ok(Self {
            config,
            parameters,
            workspace: LinearAlgebraWorkspace {
                pressure_matrix: None,
                pressure_rhs: None,
                pressure_solution: None,
                momentum_coefficients: None,
            },
            turbulence_model,
        })
    }

    /// Solve one SIMPLE iteration
    pub async fn solve_simple_iteration(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
        dt: f64,
    ) -> CfdResult<(f64, f64)> {
        // Step 1: Solve momentum equations with current pressure field
        self.momentum_prediction_step(flow_field, dt).await?;

        // Step 2: Solve pressure correction equation
        let mass_residual = self.pressure_correction_step(flow_field).await?;

        // Step 3: Correct velocities
        self.velocity_correction_step(flow_field).await?;

        // Step 4: Update pressure field
        self.pressure_update_step(flow_field).await?;

        // Step 5: Solve turbulence equations if enabled
        if self.parameters.use_turbulence {
            self.solve_turbulence(flow_field, dt).await?;
        }

        // Step 6: Apply boundary conditions
        flow_field.apply_boundary_conditions(boundary_conditions)?;

        // Step 7: Apply under-relaxation
        flow_field.apply_velocity_relaxation(self.parameters.velocity_relaxation)?;
        flow_field.apply_pressure_relaxation(self.parameters.pressure_relaxation)?;

        // Compute momentum residual
        let momentum_residual = flow_field.compute_velocity_residual();

        Ok((mass_residual, momentum_residual))
    }

    /// Momentum prediction step: solve momentum equations with current pressure
    pub async fn momentum_prediction_step(
        &self,
        flow_field: &mut FlowField,
        dt: f64,
    ) -> CfdResult<()> {
        let (_nx, _ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;
        let mu = self.config.viscosity;

        // Copy current velocities to old values for time derivatives
        flow_field.update_old_values();

        // Solve u-momentum equation
        self.solve_u_momentum(flow_field, dt, rho, mu, dx, dy)
            .await?;

        // Solve v-momentum equation
        self.solve_v_momentum(flow_field, dt, rho, mu, dx, dy)
            .await?;

        // Store predicted velocities
        flow_field.copy_to_starred();

        Ok(())
    }

    /// Solve u-momentum equation
    async fn solve_u_momentum(
        &self,
        flow_field: &mut FlowField,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();

        // For each u-velocity point (face-centered)
        for j in 1..ny - 1 {
            for i in 1..nx {
                // u goes from 1 to nx-1 for interior
                // Discretize u-momentum equation at (i, j)
                let coeffs =
                    self.compute_u_momentum_coefficients(flow_field, i, j, dt, rho, mu, dx, dy)?;

                // Solve for new u velocity using Gauss-Seidel
                let u_new = (coeffs.source
                    + coeffs.east * flow_field.u[(j, i + 1)]
                    + coeffs.west * flow_field.u[(j, i - 1)]
                    + coeffs.north * flow_field.u[(j + 1, i)]
                    + coeffs.south * flow_field.u[(j - 1, i)])
                    / coeffs.center;

                flow_field.u[(j, i)] = u_new;
            }
        }

        Ok(())
    }

    /// Solve v-momentum equation
    async fn solve_v_momentum(
        &self,
        flow_field: &mut FlowField,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();

        // For each v-velocity point (face-centered)
        for j in 1..ny {
            // v goes from 1 to ny-1 for interior
            for i in 1..nx - 1 {
                // Discretize v-momentum equation at (i, j)
                let coeffs =
                    self.compute_v_momentum_coefficients(flow_field, i, j, dt, rho, mu, dx, dy)?;

                // Solve for new v velocity using Gauss-Seidel
                let v_new = (coeffs.source
                    + coeffs.east * flow_field.v[(j, i + 1)]
                    + coeffs.west * flow_field.v[(j, i - 1)]
                    + coeffs.north * flow_field.v[(j + 1, i)]
                    + coeffs.south * flow_field.v[(j - 1, i)])
                    / coeffs.center;

                flow_field.v[(j, i)] = v_new;
            }
        }

        Ok(())
    }

    /// Pressure correction step: solve pressure Poisson equation
    pub async fn pressure_correction_step(&self, flow_field: &mut FlowField) -> CfdResult<f64> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;

        // Setup pressure correction equation: ∇²p' = ρ/Δt * ∇·u*
        for j in 1..ny - 1 {
            for i in 1..nx - 1 {
                // Compute mass source (continuity equation residual)
                let mass_source = self.compute_mass_source(flow_field, i, j, dx, dy, rho)?;
                flow_field.sp[(j, i)] = mass_source;
            }
        }

        // Solve pressure correction equation using Gauss-Seidel
        let mut max_residual = 0.0;
        for _iteration in 0..100 {
            // Inner pressure correction iterations
            let mut residual = 0.0;

            for j in 1..ny - 1 {
                for i in 1..nx - 1 {
                    let coeffs = self.compute_pressure_coefficients(dx, dy)?;

                    let p_new = (flow_field.sp[(j, i)]
                        + coeffs.east * flow_field.p_prime[(j, i + 1)]
                        + coeffs.west * flow_field.p_prime[(j, i - 1)]
                        + coeffs.north * flow_field.p_prime[(j + 1, i)]
                        + coeffs.south * flow_field.p_prime[(j - 1, i)])
                        / coeffs.center;

                    let correction = p_new - flow_field.p_prime[(j, i)];
                    residual += correction * correction;
                    flow_field.p_prime[(j, i)] = p_new;
                }
            }

            residual = residual.sqrt();
            max_residual = residual;

            if residual < 1e-8 {
                break;
            }
        }

        Ok(max_residual)
    }

    /// Velocity correction step: correct velocities with pressure correction
    pub async fn velocity_correction_step(&self, flow_field: &mut FlowField) -> CfdResult<()> {
        let (nx, ny, dx, dy) = flow_field.grid_info();
        let rho = self.config.density;

        // Correct u-velocities
        for j in 1..ny - 1 {
            for i in 1..nx {
                if i > 0 && i < nx {
                    let dp_dx = (flow_field.p_prime[(j, i)] - flow_field.p_prime[(j, i - 1)]) / dx;
                    let ap_u =
                        self.compute_u_momentum_center_coefficient(flow_field, i, j, dx, dy, rho)?;
                    flow_field.u[(j, i)] = flow_field.u_star[(j, i)] - (dx * dy / ap_u) * dp_dx;
                }
            }
        }

        // Correct v-velocities
        for j in 1..ny {
            for i in 1..nx - 1 {
                if j > 0 && j < ny {
                    let dp_dy = (flow_field.p_prime[(j, i)] - flow_field.p_prime[(j - 1, i)]) / dy;
                    let ap_v =
                        self.compute_v_momentum_center_coefficient(flow_field, i, j, dx, dy, rho)?;
                    flow_field.v[(j, i)] = flow_field.v_star[(j, i)] - (dx * dy / ap_v) * dp_dy;
                }
            }
        }

        Ok(())
    }

    /// Pressure update step: add pressure correction to pressure
    pub async fn pressure_update_step(&self, flow_field: &mut FlowField) -> CfdResult<()> {
        let (nx, ny, _, _) = flow_field.grid_info();

        for j in 0..ny {
            for i in 0..nx {
                flow_field.p[(j, i)] +=
                    self.parameters.pressure_relaxation * flow_field.p_prime[(j, i)];
                flow_field.p_prime[(j, i)] = 0.0; // Reset pressure correction
            }
        }

        Ok(())
    }

    /// Solve turbulence model equations
    async fn solve_turbulence(&mut self, flow_field: &FlowField, dt: f64) -> CfdResult<()> {
        if self.parameters.use_turbulence {
            // Initialize turbulence model if not already done
            let (nx, ny, dx, dy) = flow_field.grid_info();
            let n_cells = nx * ny;

            if self.turbulence_model.is_none() {
                self.turbulence_model =
                    Some(KEpsilonModel::new(KEpsilonVariant::Standard, n_cells));
            }

            let turbulence_model = self.turbulence_model.as_mut().unwrap();

            // Convert velocity field to Vector3 format
            let mut velocity_vec = Vec::with_capacity(n_cells);
            for j in 0..ny {
                for i in 0..nx {
                    let u = if i < nx && j < ny {
                        flow_field.u[(j, i)]
                    } else {
                        0.0
                    };
                    let v = if i < nx && j < ny {
                        flow_field.v[(j, i)]
                    } else {
                        0.0
                    };
                    velocity_vec.push(Vector3::new(u, v, 0.0)); // 2D case, w=0
                }
            }

            // Simplified velocity gradients (zero for now)
            let velocity_gradients = vec![[[0.0; 3]; 3]; n_cells];

            // Create pressure vector
            let mut pressure = DVector::zeros(n_cells);
            for j in 0..ny {
                for i in 0..nx {
                    if i < nx && j < ny {
                        pressure[j * nx + i] = flow_field.p[(j, i)];
                    }
                }
            }

            let turbulence_state = TurbulenceState {
                velocity: velocity_vec,
                velocity_gradients,
                pressure,
                turbulent_ke: None, // Will be initialized by model
                epsilon: None,      // Will be initialized by model
                omega: None,
                wall_distance: DVector::from_element(n_cells, 1.0), // Simplified
                cell_volumes: DVector::from_element(n_cells, dx * dy), // 2D cell volume
                molecular_viscosity: self.config.viscosity / self.config.density, // kinematic viscosity
                density: self.config.density,
            };

            // Update turbulence model with new state
            turbulence_model.update(&turbulence_state, dt)?;
        }
        Ok(())
    }

    /// Compute effective viscosity (molecular + turbulent)
    fn compute_effective_viscosity(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        mu: f64,
    ) -> f64 {
        if let Some(ref turbulence_model) = self.turbulence_model {
            // Get turbulent viscosity from the model
            let (nx, _ny, _, _) = flow_field.grid_info();
            let cell_idx = j * nx + i;
            if cell_idx < turbulence_model.nu_t.len() {
                let nu_t = turbulence_model.nu_t[cell_idx]; // kinematic turbulent viscosity
                let mu_t = nu_t * self.config.density; // convert to dynamic viscosity
                mu + mu_t
            } else {
                mu
            }
        } else {
            mu
        }
    }

    /// Compute coefficients for u-momentum equation
    fn compute_u_momentum_coefficients(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<MomentumEquationCoeffs> {
        // Compute effective viscosity (molecular + turbulent)
        let mu_eff = self.compute_effective_viscosity(flow_field, i, j, mu);

        // Diffusion coefficients using effective viscosity
        let gamma_e = mu_eff / dx;
        let gamma_w = mu_eff / dx;
        let gamma_n = mu_eff / dy;
        let gamma_s = mu_eff / dy;

        // Convection coefficients (using upwind)
        let (u_center, v_center) = flow_field.get_velocity_at(i, j)?;
        let fe = rho * u_center * dy; // East face mass flux
        let fw = rho * u_center * dy; // West face mass flux
        let fn_ = rho * v_center * dx; // North face mass flux
        let fs = rho * v_center * dx; // South face mass flux

        // Compute coefficients with upwind scheme
        let ae = gamma_e + f64::max(-fe, 0.0);
        let aw = gamma_w + f64::max(fw, 0.0);
        let an = gamma_n + f64::max(-fn_, 0.0);
        let as_ = gamma_s + f64::max(fs, 0.0);

        // Time derivative coefficient
        let ap0 = rho * dx * dy / dt;

        // Central coefficient
        let ap = ae + aw + an + as_ + ap0;

        // Source term (pressure gradient + old time step)
        let pressure_gradient = -(flow_field.p[(j, i)] - flow_field.p[(j, i - 1)]) * dy;
        let time_term = ap0 * flow_field.u_old[(j, i)];
        let source = pressure_gradient + time_term;

        Ok(MomentumEquationCoeffs {
            center: ap,
            east: ae,
            west: aw,
            north: an,
            south: as_,
            source,
        })
    }

    /// Compute coefficients for v-momentum equation
    fn compute_v_momentum_coefficients(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dt: f64,
        rho: f64,
        mu: f64,
        dx: f64,
        dy: f64,
    ) -> CfdResult<MomentumEquationCoeffs> {
        // Compute effective viscosity (molecular + turbulent)
        let mu_eff = self.compute_effective_viscosity(flow_field, i, j, mu);

        // Similar to u-momentum but for v-component
        let gamma_e = mu_eff / dx;
        let gamma_w = mu_eff / dx;
        let gamma_n = mu_eff / dy;
        let gamma_s = mu_eff / dy;

        let (u_center, v_center) = flow_field.get_velocity_at(i, j)?;
        let fe = rho * u_center * dy;
        let fw = rho * u_center * dy;
        let fn_ = rho * v_center * dx;
        let fs = rho * v_center * dx;

        let ae = gamma_e + f64::max(-fe, 0.0);
        let aw = gamma_w + f64::max(fw, 0.0);
        let an = gamma_n + f64::max(-fn_, 0.0);
        let as_ = gamma_s + f64::max(fs, 0.0);

        let ap0 = rho * dx * dy / dt;
        let ap = ae + aw + an + as_ + ap0;

        // Pressure gradient in y-direction
        let pressure_gradient = -(flow_field.p[(j, i)] - flow_field.p[(j - 1, i)]) * dx;
        let time_term = ap0 * flow_field.v_old[(j, i)];
        let source = pressure_gradient + time_term;

        Ok(MomentumEquationCoeffs {
            center: ap,
            east: ae,
            west: aw,
            north: an,
            south: as_,
            source,
        })
    }

    /// Compute mass source term for pressure correction equation
    fn compute_mass_source(
        &self,
        flow_field: &FlowField,
        i: usize,
        j: usize,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<f64> {
        // Mass source = ρ * ∇·u*
        let u_e = flow_field.u_star[(j, i + 1)];
        let u_w = flow_field.u_star[(j, i)];
        let v_n = flow_field.v_star[(j + 1, i)];
        let v_s = flow_field.v_star[(j, i)];

        let mass_flux_imbalance = rho * ((u_e - u_w) * dy + (v_n - v_s) * dx);

        Ok(-mass_flux_imbalance) // Negative because we want ∇²p' = -∇·u*
    }

    /// Compute coefficients for pressure correction equation
    fn compute_pressure_coefficients(&self, dx: f64, dy: f64) -> CfdResult<MomentumEquationCoeffs> {
        // Pressure correction equation: ∇²p' = S
        // Standard 5-point stencil with unit coefficients
        let ae = 1.0 / (dx * dx);
        let aw = 1.0 / (dx * dx);
        let an = 1.0 / (dy * dy);
        let as_ = 1.0 / (dy * dy);
        let ap = ae + aw + an + as_;

        Ok(MomentumEquationCoeffs {
            center: ap,
            east: ae,
            west: aw,
            north: an,
            south: as_,
            source: 0.0, // Source is set separately
        })
    }

    /// Compute center coefficient for u-momentum equation
    fn compute_u_momentum_center_coefficient(
        &self,
        _flow_field: &FlowField,
        _i: usize,
        _j: usize,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<f64> {
        // Simplified calculation for velocity correction
        // This would be the diagonal coefficient from momentum discretization
        Ok(rho * dx * dy / self.parameters.time_step)
    }

    /// Compute center coefficient for v-momentum equation
    fn compute_v_momentum_center_coefficient(
        &self,
        _flow_field: &FlowField,
        _i: usize,
        _j: usize,
        dx: f64,
        dy: f64,
        rho: f64,
    ) -> CfdResult<f64> {
        Ok(rho * dx * dy / self.parameters.time_step)
    }
}

/// Coefficients for momentum equation discretization
#[derive(Debug, Clone)]
struct MomentumEquationCoeffs {
    pub center: f64,
    pub east: f64,
    pub west: f64,
    pub north: f64,
    pub south: f64,
    pub source: f64,
}

#[async_trait]
impl IncompressibleSolver for SimpleSolver {
    type Parameters = SimpleParameters;
    type Result = SimpleResult;

    fn new(config: CfdConfig, params: Self::Parameters) -> CfdResult<Self> {
        Self::new(config, params)
    }

    async fn solve_time_step(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
        dt: f64,
    ) -> CfdResult<Self::Result> {
        let start_time = Instant::now();
        let mut residual_history = Vec::new();
        let pressure_iterations = Vec::new();

        for iteration in 0..self.parameters.max_iterations {
            let (mass_residual, momentum_residual) = self
                .solve_simple_iteration(flow_field, boundary_conditions, dt)
                .await?;

            let total_residual =
                (mass_residual * mass_residual + momentum_residual * momentum_residual).sqrt();
            residual_history.push(total_residual);

            if total_residual < self.parameters.tolerance {
                let solve_time = start_time.elapsed();

                return Ok(SimpleResult {
                    solver_result: SolverResult {
                        converged: true,
                        iterations: iteration + 1,
                        final_residual: total_residual,
                        residual_history,
                        solve_time,
                    },
                    pressure_iterations,
                    mass_residual,
                    momentum_residual,
                });
            }
        }

        // Did not converge
        let solve_time = start_time.elapsed();
        Ok(SimpleResult {
            solver_result: SolverResult {
                converged: false,
                iterations: self.parameters.max_iterations,
                final_residual: residual_history.last().copied().unwrap_or(f64::INFINITY),
                residual_history,
                solve_time,
            },
            pressure_iterations,
            mass_residual: f64::INFINITY,
            momentum_residual: f64::INFINITY,
        })
    }

    async fn solve(
        &mut self,
        flow_field: &mut FlowField,
        boundary_conditions: &BoundaryConditions,
    ) -> CfdResult<Self::Result> {
        // For steady-state solve, use default time step
        self.solve_time_step(flow_field, boundary_conditions, self.parameters.time_step)
            .await
    }

    fn config(&self) -> &CfdConfig {
        &self.config
    }

    fn parameters(&self) -> &Self::Parameters {
        &self.parameters
    }
}