//! Flow field data structures and operations
//!
//! This module defines the core data structures for storing and manipulating
//! flow field variables (velocity, pressure) on structured grids.

use crate::{CfdError, CfdResult};
use nalgebra::DMatrix;

/// Flow field containing all flow variables on a structured grid
#[derive(Debug, Clone)]
pub struct FlowField {
    /// Grid dimensions
    pub nx: usize,
    pub ny: usize,

    /// Grid spacing
    pub dx: f64,
    pub dy: f64,

    /// x-component of velocity (u) - located at cell faces (i+1/2, j)
    pub u: DMatrix<f64>,
    /// y-component of velocity (v) - located at cell faces (i, j+1/2)
    pub v: DMatrix<f64>,
    /// Pressure (p) - located at cell centers (i, j)
    pub p: DMatrix<f64>,

    /// Previous time step values for time integration
    pub u_old: DMatrix<f64>,
    pub v_old: DMatrix<f64>,
    pub p_old: DMatrix<f64>,

    /// Auxiliary fields for solver algorithms
    pub u_star: DMatrix<f64>, // Predicted velocity (SIMPLE/PISO)
    pub v_star: DMatrix<f64>,  // Predicted velocity (SIMPLE/PISO)
    pub p_prime: DMatrix<f64>, // Pressure correction (SIMPLE/PISO)

    /// Source terms
    pub su: DMatrix<f64>, // u-momentum source
    pub sv: DMatrix<f64>, // v-momentum source
    pub sp: DMatrix<f64>, // Pressure source (mass source)
}

impl FlowField {
    /// Create new flow field with given dimensions
    pub fn new(nx: usize, ny: usize, dx: f64, dy: f64) -> CfdResult<Self> {
        if nx < 3 || ny < 3 {
            return Err(CfdError::invalid_parameter("Grid must be at least 3x3"));
        }

        if dx <= 0.0 || dy <= 0.0 {
            return Err(CfdError::invalid_parameter("Grid spacing must be positive"));
        }

        // For staggered grid:
        // u: (nx+1, ny) - face-centered in x-direction
        // v: (nx, ny+1) - face-centered in y-direction
        // p: (nx, ny) - cell-centered

        let zeros_u = DMatrix::zeros(ny, nx + 1); // Note: nalgebra is (rows, cols)
        let zeros_v = DMatrix::zeros(ny + 1, nx);
        let zeros_p = DMatrix::zeros(ny, nx);

        Ok(Self {
            nx,
            ny,
            dx,
            dy,

            u: zeros_u.clone(),
            v: zeros_v.clone(),
            p: zeros_p.clone(),

            u_old: zeros_u.clone(),
            v_old: zeros_v.clone(),
            p_old: zeros_p.clone(),

            u_star: zeros_u.clone(),
            v_star: zeros_v.clone(),
            p_prime: zeros_p.clone(),

            su: zeros_u,
            sv: zeros_v,
            sp: zeros_p,
        })
    }

    /// Set velocity at a given grid point
    pub fn set_velocity(&mut self, i: usize, j: usize, u_val: f64, v_val: f64) -> CfdResult<()> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        // For staggered grid, velocity components are at different locations
        // u is stored at (j, i) for face (i+1/2, j)
        if i < self.nx {
            self.u[(j, i)] = u_val;
        }

        // v is stored at (j, i) for face (i, j+1/2)
        if j < self.ny {
            self.v[(j, i)] = v_val;
        }

        Ok(())
    }

    /// Get velocity at a given grid point (interpolated to cell center)
    pub fn get_velocity_at(&self, i: usize, j: usize) -> CfdResult<(f64, f64)> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        // Interpolate velocities to cell center
        let u_center = if i == 0 {
            self.u[(j, 0)]
        } else if i == self.nx - 1 {
            self.u[(j, self.nx - 1)]
        } else {
            0.5 * (self.u[(j, i - 1)] + self.u[(j, i)])
        };

        let v_center = if j == 0 {
            self.v[(0, i)]
        } else if j == self.ny - 1 {
            self.v[(self.ny - 1, i)]
        } else {
            0.5 * (self.v[(j - 1, i)] + self.v[(j, i)])
        };

        Ok((u_center, v_center))
    }

    /// Set pressure at a given grid point
    pub fn set_pressure(&mut self, i: usize, j: usize, p_val: f64) -> CfdResult<()> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        self.p[(j, i)] = p_val;
        Ok(())
    }

    /// Get pressure at a given grid point
    pub fn get_pressure_at(&self, i: usize, j: usize) -> CfdResult<f64> {
        if i >= self.nx || j >= self.ny {
            return Err(CfdError::invalid_parameter("Grid indices out of bounds"));
        }

        Ok(self.p[(j, i)])
    }

    /// Apply boundary conditions to the flow field
    pub fn apply_boundary_conditions(&mut self, bcs: &super::BoundaryConditions) -> CfdResult<()> {
        bcs.apply_to_flow_field(self)
    }

    /// Compute divergence of velocity field (mass conservation check)
    pub fn compute_divergence(&self) -> CfdResult<DMatrix<f64>> {
        let mut divergence = DMatrix::zeros(self.ny, self.nx);

        for j in 0..self.ny {
            for i in 0..self.nx {
                // ∇·u = ∂u/∂x + ∂v/∂y
                let du_dx = if i == self.nx - 1 {
                    (self.u[(j, i)] - self.u[(j, i - 1)]) / self.dx
                } else {
                    (self.u[(j, i + 1)] - self.u[(j, i)]) / self.dx
                };

                let dv_dy = if j == self.ny - 1 {
                    (self.v[(j, i)] - self.v[(j - 1, i)]) / self.dy
                } else {
                    (self.v[(j + 1, i)] - self.v[(j, i)]) / self.dy
                };

                divergence[(j, i)] = du_dx + dv_dy;
            }
        }

        Ok(divergence)
    }

    /// Compute maximum divergence (for mass conservation check)
    pub fn compute_max_divergence(&self) -> CfdResult<f64> {
        let divergence = self.compute_divergence()?;
        Ok(divergence.iter().map(|&x| x.abs()).fold(0.0, f64::max))
    }

    /// Find maximum u-velocity and its location
    pub fn find_max_u_velocity(&self) -> CfdResult<(f64, (usize, usize))> {
        let mut max_u = f64::NEG_INFINITY;
        let mut max_loc = (0, 0);

        for j in 0..self.ny {
            for i in 0..=self.nx {
                if self.u[(j, i)] > max_u {
                    max_u = self.u[(j, i)];
                    max_loc = (i, j);
                }
            }
        }

        Ok((max_u, max_loc))
    }

    /// Compute total kinetic energy
    pub fn compute_kinetic_energy(&self) -> CfdResult<f64> {
        let mut ke = 0.0;

        for j in 0..self.ny {
            for i in 0..self.nx {
                let (u_center, v_center) = self.get_velocity_at(i, j)?;
                ke += 0.5 * (u_center * u_center + v_center * v_center) * self.dx * self.dy;
            }
        }

        Ok(ke)
    }

    /// Update old values (for time stepping)
    pub fn update_old_values(&mut self) {
        self.u_old.copy_from(&self.u);
        self.v_old.copy_from(&self.v);
        self.p_old.copy_from(&self.p);
    }

    /// Copy current values to starred values (for predictor step)
    pub fn copy_to_starred(&mut self) {
        self.u_star.copy_from(&self.u);
        self.v_star.copy_from(&self.v);
    }

    /// Apply under-relaxation to velocity field
    pub fn apply_velocity_relaxation(&mut self, relaxation_factor: f64) -> CfdResult<()> {
        if relaxation_factor <= 0.0 || relaxation_factor > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Relaxation factor must be in (0, 1]",
            ));
        }

        // u = α * u_new + (1 - α) * u_old
        for j in 0..self.ny {
            for i in 0..=self.nx {
                self.u[(j, i)] = relaxation_factor * self.u[(j, i)]
                    + (1.0 - relaxation_factor) * self.u_old[(j, i)];
            }
        }

        for j in 0..=self.ny {
            for i in 0..self.nx {
                self.v[(j, i)] = relaxation_factor * self.v[(j, i)]
                    + (1.0 - relaxation_factor) * self.v_old[(j, i)];
            }
        }

        Ok(())
    }

    /// Apply under-relaxation to pressure field
    pub fn apply_pressure_relaxation(&mut self, relaxation_factor: f64) -> CfdResult<()> {
        if relaxation_factor <= 0.0 || relaxation_factor > 1.0 {
            return Err(CfdError::invalid_parameter(
                "Relaxation factor must be in (0, 1]",
            ));
        }

        for j in 0..self.ny {
            for i in 0..self.nx {
                self.p[(j, i)] = relaxation_factor * self.p[(j, i)]
                    + (1.0 - relaxation_factor) * self.p_old[(j, i)];
            }
        }

        Ok(())
    }

    /// Compute L2 norm of residual
    #[must_use]
    pub fn compute_velocity_residual(&self) -> f64 {
        let mut residual = 0.0;

        // u-momentum residual
        for j in 0..self.ny {
            for i in 0..=self.nx {
                let diff = self.u[(j, i)] - self.u_old[(j, i)];
                residual += diff * diff;
            }
        }

        // v-momentum residual
        for j in 0..=self.ny {
            for i in 0..self.nx {
                let diff = self.v[(j, i)] - self.v_old[(j, i)];
                residual += diff * diff;
            }
        }

        residual.sqrt()
    }

    /// Compute pressure residual
    #[must_use]
    pub fn compute_pressure_residual(&self) -> f64 {
        let mut residual = 0.0;

        for j in 0..self.ny {
            for i in 0..self.nx {
                let diff = self.p[(j, i)] - self.p_old[(j, i)];
                residual += diff * diff;
            }
        }

        residual.sqrt()
    }

    /// Get grid information
    #[must_use]
    pub fn grid_info(&self) -> (usize, usize, f64, f64) {
        (self.nx, self.ny, self.dx, self.dy)
    }

    /// Initialize with analytical solution (for testing)
    pub fn initialize_with_analytical(
        &mut self,
        solution_type: AnalyticalSolution,
    ) -> CfdResult<()> {
        match solution_type {
            AnalyticalSolution::PoiseuillePlane { u_max } => {
                // Plane Poiseuille flow: u(y) = u_max * 4 * y * (1-y)
                for j in 0..self.ny {
                    let y = (j as f64 + 0.5) * self.dy; // Cell center y-coordinate
                    let y_normalized = y / (self.ny as f64 * self.dy);
                    let u_analytical = u_max * 4.0 * y_normalized * (1.0 - y_normalized);

                    for i in 0..=self.nx {
                        self.u[(j, i)] = u_analytical;
                    }
                }

                // v = 0 everywhere
                self.v.fill(0.0);

                // Pressure gradient to drive the flow
                for j in 0..self.ny {
                    for i in 0..self.nx {
                        self.p[(j, i)] = -(i as f64) * self.dx; // Linear pressure drop
                    }
                }
            }
            AnalyticalSolution::TaylorGreenVortex { amplitude } => {
                // Taylor-Green vortex: analytical solution for 2D Navier-Stokes
                for j in 0..self.ny {
                    for i in 0..=self.nx {
                        let x = i as f64 * self.dx;
                        let y = (j as f64 + 0.5) * self.dy;
                        self.u[(j, i)] = amplitude
                            * (2.0 * std::f64::consts::PI * x).sin()
                            * (2.0 * std::f64::consts::PI * y).cos();
                    }
                }

                for j in 0..=self.ny {
                    for i in 0..self.nx {
                        let x = (i as f64 + 0.5) * self.dx;
                        let y = j as f64 * self.dy;
                        self.v[(j, i)] = -amplitude
                            * (2.0 * std::f64::consts::PI * x).cos()
                            * (2.0 * std::f64::consts::PI * y).sin();
                    }
                }

                // Pressure field for Taylor-Green vortex
                for j in 0..self.ny {
                    for i in 0..self.nx {
                        let x = (i as f64 + 0.5) * self.dx;
                        let y = (j as f64 + 0.5) * self.dy;
                        self.p[(j, i)] = -0.25
                            * amplitude
                            * amplitude
                            * ((4.0 * std::f64::consts::PI * x).cos()
                                + (4.0 * std::f64::consts::PI * y).cos());
                    }
                }
            }
        }

        Ok(())
    }
}

/// Analytical solutions for testing and validation
#[derive(Debug, Clone, Copy)]
pub enum AnalyticalSolution {
    /// Plane Poiseuille flow between parallel plates
    PoiseuillePlane { u_max: f64 },
    /// Taylor-Green vortex (decaying vortex solution)
    TaylorGreenVortex { amplitude: f64 },
}