//! D2Q9 Lattice Boltzmann Method implementation
//!
//! This module implements the D2Q9 (2D, 9 velocities) lattice Boltzmann method
//! for simulating incompressible fluid flows.

use crate::error::{CfdError, CfdResult};
use crate::solvers::lbm::common::MacroscopicVariables;
use nalgebra::Vector2;

/// D2Q9 lattice velocities (in lattice units)
const D2Q9_VELOCITIES: [Vector2<i32>; 9] = [
    Vector2::new(0, 0),   // 0: Rest particle
    Vector2::new(1, 0),   // 1: East
    Vector2::new(0, 1),   // 2: North
    Vector2::new(-1, 0),  // 3: West
    Vector2::new(0, -1),  // 4: South
    Vector2::new(1, 1),   // 5: Northeast
    Vector2::new(-1, 1),  // 6: Northwest
    Vector2::new(-1, -1), // 7: Southwest
    Vector2::new(1, -1),  // 8: Southeast
];

/// D2Q9 lattice weights
const D2Q9_WEIGHTS: [f64; 9] = [
    4.0 / 9.0, // 0: Rest particle
    1.0 / 9.0, // 1-4: Cardinal directions
    1.0 / 9.0,
    1.0 / 9.0,
    1.0 / 9.0,
    1.0 / 36.0, // 5-8: Diagonal directions
    1.0 / 36.0,
    1.0 / 36.0,
    1.0 / 36.0,
];

/// Parameters for D2Q9 lattice Boltzmann method
#[derive(Debug, Clone)]
pub struct D2Q9Parameters {
    /// Relaxation time for BGK collision operator
    pub tau: f64,
    /// Kinematic viscosity (derived from tau)
    pub nu: f64,
}

impl D2Q9Parameters {
    /// Create new D2Q9 parameters with given relaxation time
    #[must_use]
    pub fn new(tau: f64) -> Self {
        let nu = (tau - 0.5) / 3.0; // Relationship between tau and viscosity
        Self { tau, nu }
    }

    /// Validate parameters for numerical stability
    pub fn validate(&self) -> CfdResult<()> {
        if self.tau <= 0.5 {
            return Err(CfdError::invalid_parameter(
                "Relaxation time must be greater than 0.5 for stability",
            ));
        }
        if self.tau >= 2.0 {
            return Err(CfdError::invalid_parameter(
                "Relaxation time should be less than 2.0 for efficiency",
            ));
        }
        if self.nu <= 0.0 {
            return Err(CfdError::invalid_parameter(
                "Kinematic viscosity must be positive",
            ));
        }
        Ok(())
    }

    /// Create parameters from kinematic viscosity
    #[must_use]
    pub fn from_viscosity(nu: f64) -> Self {
        let tau = 3.0 * nu + 0.5;
        Self { tau, nu }
    }
}

impl Default for D2Q9Parameters {
    fn default() -> Self {
        Self::new(0.6) // Commonly used value
    }
}

/// D2Q9 Lattice Boltzmann Method solver
pub struct D2Q9Solver {
    /// Grid dimensions
    nx: usize,
    ny: usize,

    /// Distribution functions f[x][y][i] where i is velocity direction
    f: Vec<Vec<Vec<f64>>>,

    /// Temporary storage for streaming step
    f_temp: Vec<Vec<Vec<f64>>>,

    /// Solver parameters
    params: D2Q9Parameters,
}

impl D2Q9Solver {
    /// Create new D2Q9 solver
    #[must_use]
    pub fn new(nx: usize, ny: usize, params: D2Q9Parameters) -> Self {
        params.validate().expect("Invalid D2Q9 parameters");

        let f = vec![vec![vec![0.0; 9]; ny]; nx];
        let f_temp = vec![vec![vec![0.0; 9]; ny]; nx];

        Self {
            nx,
            ny,
            f,
            f_temp,
            params,
        }
    }

    /// Get lattice velocities
    #[must_use]
    pub fn lattice_velocities(&self) -> Vec<Vector2<i32>> {
        D2Q9_VELOCITIES.to_vec()
    }

    /// Get lattice weights
    #[must_use]
    pub fn weights(&self) -> Vec<f64> {
        D2Q9_WEIGHTS.to_vec()
    }

    /// Calculate equilibrium distribution function
    #[must_use]
    pub fn equilibrium_distribution(&self, density: f64, velocity: &Vector2<f64>) -> Vec<f64> {
        let mut f_eq = vec![0.0; 9];
        let u_sqr = velocity.norm_squared();

        for i in 0..9 {
            let e_i = Vector2::new(
                f64::from(D2Q9_VELOCITIES[i].x),
                f64::from(D2Q9_VELOCITIES[i].y),
            );
            let e_dot_u = e_i.dot(velocity);

            // Equilibrium distribution: f_i^eq = w_i * rho * (1 + 3*e_i·u + 9/2*(e_i·u)^2 - 3/2*u^2)
            f_eq[i] = D2Q9_WEIGHTS[i]
                * density
                * (1.0 + 3.0 * e_dot_u + 4.5 * e_dot_u * e_dot_u - 1.5 * u_sqr);
        }

        f_eq
    }

    /// Set distribution function at a specific grid point
    pub fn set_distribution_at(&mut self, x: usize, y: usize, f_values: &[f64]) {
        assert_eq!(f_values.len(), 9);
        self.f[x][y].copy_from_slice(f_values);
    }

    /// Get distribution function at a specific grid point
    #[must_use]
    pub fn distribution_at(&self, x: usize, y: usize) -> Vec<f64> {
        self.f[x][y].clone()
    }

    /// Extract macroscopic variables (density and velocity) from distribution functions
    #[must_use]
    pub fn macroscopic_variables_at(&self, x: usize, y: usize) -> MacroscopicVariables {
        let f_local = &self.f[x][y];

        // Density: sum of all distribution functions
        let density: f64 = f_local.iter().sum();

        // Momentum: sum of f_i * e_i
        let mut momentum = Vector2::zeros();
        for i in 0..9 {
            let e_i = Vector2::new(
                f64::from(D2Q9_VELOCITIES[i].x),
                f64::from(D2Q9_VELOCITIES[i].y),
            );
            momentum += f_local[i] * e_i;
        }

        // Velocity: momentum / density
        let velocity = if density > 1e-15 {
            momentum / density
        } else {
            Vector2::zeros()
        };

        MacroscopicVariables::new(density, velocity)
    }

    /// BGK collision step
    pub fn collision_step(&mut self) {
        let omega = 1.0 / self.params.tau; // Collision frequency

        for x in 0..self.nx {
            for y in 0..self.ny {
                let vars = self.macroscopic_variables_at(x, y);
                let f_eq = self.equilibrium_distribution(vars.density, &vars.velocity);

                // BGK collision: f_i^new = f_i - omega * (f_i - f_i^eq)
                for i in 0..9 {
                    self.f[x][y][i] -= omega * (self.f[x][y][i] - f_eq[i]);
                }
            }
        }
    }

    /// Streaming step (propagation)
    pub fn streaming_step(&mut self) {
        // Copy current state to temporary storage
        for x in 0..self.nx {
            for y in 0..self.ny {
                self.f_temp[x][y].copy_from_slice(&self.f[x][y]);
            }
        }

        // Stream particles according to their velocities
        for x in 0..self.nx {
            for y in 0..self.ny {
                for i in 0..9 {
                    let e_i = D2Q9_VELOCITIES[i];
                    let x_src = (x as i32 - e_i.x).rem_euclid(self.nx as i32) as usize;
                    let y_src = (y as i32 - e_i.y).rem_euclid(self.ny as i32) as usize;

                    self.f[x][y][i] = self.f_temp[x_src][y_src][i];
                }
            }
        }
    }

    /// Complete LBM time step (collision + streaming)
    pub fn step(&mut self) {
        self.collision_step();
        self.streaming_step();
        self.apply_bounce_back_boundaries();
    }

    /// Complete LBM time step with specified boundary condition function
    pub fn step_with_boundaries<F>(&mut self, apply_boundaries: F)
    where
        F: FnOnce(&mut Self),
    {
        self.collision_step();
        self.streaming_step();
        apply_boundaries(self);
    }

    /// Complete LBM time step without boundary conditions (for periodic domains)
    pub fn step_periodic(&mut self) {
        self.collision_step();
        self.streaming_step();
    }

    /// Initialize uniform flow field
    pub fn initialize_uniform(&mut self, density: f64, velocity: Vector2<f64>) {
        let f_eq = self.equilibrium_distribution(density, &velocity);

        for x in 0..self.nx {
            for y in 0..self.ny {
                self.f[x][y].copy_from_slice(&f_eq);
            }
        }
    }

    /// Initialize Poiseuille flow (parabolic velocity profile)
    pub fn initialize_poiseuille_flow(&mut self, driving_force: f64) {
        let density = 1.0;

        for x in 0..self.nx {
            for y in 0..self.ny {
                if y == 0 || y == self.ny - 1 {
                    // No-slip boundary conditions at walls
                    let f_eq = self.equilibrium_distribution(density, &Vector2::zeros());
                    self.f[x][y].copy_from_slice(&f_eq);
                } else {
                    // Parabolic velocity profile for interior points
                    let y_normalized = y as f64 / (self.ny - 1) as f64;
                    let u_x = driving_force * 4.0 * y_normalized * (1.0 - y_normalized);
                    let velocity = Vector2::new(u_x, 0.0);

                    let f_eq = self.equilibrium_distribution(density, &velocity);
                    self.f[x][y].copy_from_slice(&f_eq);
                }
            }
        }
    }

    /// Apply bounce-back boundary conditions for no-slip walls (mass conserving)
    pub fn apply_bounce_back_boundaries(&mut self) {
        // Full-way bounce-back on the top and bottom walls: at a wall node,
        // each population is exchanged with the one travelling in the exactly
        // opposite direction. In D2Q9 those pairs are 1<->3, 2<->4, 5<->7 and
        // 6<->8.
        //
        // Two things were wrong here.
        //
        // First, this was written as assignment (`f[2] = f[4]`) rather than a
        // swap, so the population being reflected was overwritten and its
        // value discarded. Bounce-back is a permutation of populations and
        // conserves mass exactly; assignment leaks it. Over 100 steps on a
        // 16x16 lattice the domain lost 0.013% of its mass, and it would keep
        // draining for as long as the simulation ran.
        //
        // Second, the pairs used were 5<->8 and 6<->7, which reverse only the
        // wall-normal component. That is *specular* reflection, a free-slip
        // wall. Full bounce-back reverses both components and is what gives
        // the no-slip condition the walls are supposed to impose, so the
        // tangential velocity never went to zero at the wall.
        const OPPOSITE_PAIRS: [(usize, usize); 4] = [(1, 3), (2, 4), (5, 7), (6, 8)];

        let top_y = self.ny - 1;
        for x in 0..self.nx {
            for (a, b) in OPPOSITE_PAIRS {
                self.f[x][0].swap(a, b);
                self.f[x][top_y].swap(a, b);
            }
        }
    }

    /// Apply no-slip boundary conditions for walls (for initialization only)
    pub fn apply_no_slip_boundaries(&mut self) {
        let density = 1.0;
        let zero_velocity = Vector2::zeros();
        let f_eq = self.equilibrium_distribution(density, &zero_velocity);

        // Bottom and top walls
        for x in 0..self.nx {
            self.f[x][0].copy_from_slice(&f_eq);
            self.f[x][self.ny - 1].copy_from_slice(&f_eq);
        }
    }

    /// Calculate total mass in the domain
    #[must_use]
    pub fn total_mass(&self) -> f64 {
        let mut total = 0.0;
        for x in 0..self.nx {
            for y in 0..self.ny {
                let vars = self.macroscopic_variables_at(x, y);
                total += vars.density;
            }
        }
        total
    }

    /// Get grid dimensions
    #[must_use]
    pub fn dimensions(&self) -> (usize, usize) {
        (self.nx, self.ny)
    }

    /// Get solver parameters
    #[must_use]
    pub fn parameters(&self) -> &D2Q9Parameters {
        &self.params
    }

    /// Calculate kinetic energy in the domain
    #[must_use]
    pub fn kinetic_energy(&self) -> f64 {
        let mut total_ke = 0.0;
        for x in 0..self.nx {
            for y in 0..self.ny {
                let vars = self.macroscopic_variables_at(x, y);
                total_ke += 0.5 * vars.density * vars.velocity.norm_squared();
            }
        }
        total_ke
    }

    /// Calculate maximum velocity in the domain
    #[must_use]
    pub fn max_velocity(&self) -> f64 {
        let mut max_vel = 0.0f64;
        for x in 0..self.nx {
            for y in 0..self.ny {
                let vars = self.macroscopic_variables_at(x, y);
                max_vel = max_vel.max(vars.velocity.norm());
            }
        }
        max_vel
    }

    /// Check CFL condition for numerical stability
    #[must_use]
    pub fn check_cfl_condition(&self) -> bool {
        let max_vel = self.max_velocity();
        // CFL condition: max_velocity * dt / dx < 1
        // In LBM, dt = dx = 1 in lattice units, so we need max_vel < 1
        max_vel < 0.1 // Conservative limit
    }
}

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

    #[test]
    fn test_d2q9_parameters() {
        let params = D2Q9Parameters::new(0.6);
        assert_relative_eq!(params.tau, 0.6);
        assert_relative_eq!(params.nu, (0.6 - 0.5) / 3.0);
    }

    #[test]
    fn test_d2q9_from_viscosity() {
        let nu = 0.1;
        let params = D2Q9Parameters::from_viscosity(nu);
        assert_relative_eq!(params.nu, nu);
        assert_relative_eq!(params.tau, 3.0 * nu + 0.5);
    }

    #[test]
    fn test_d2q9_grid_creation() {
        let solver = D2Q9Solver::new(10, 8, D2Q9Parameters::default());
        assert_eq!(solver.dimensions(), (10, 8));
    }

    #[test]
    fn test_equilibrium_mass_conservation() {
        let solver = D2Q9Solver::new(5, 5, D2Q9Parameters::default());
        let density = 1.5;
        let velocity = Vector2::new(0.1, -0.05);

        let f_eq = solver.equilibrium_distribution(density, &velocity);
        let sum: f64 = f_eq.iter().sum();

        assert_relative_eq!(sum, density, epsilon = 1e-15);
    }

    #[test]
    fn test_equilibrium_momentum_conservation() {
        let solver = D2Q9Solver::new(5, 5, D2Q9Parameters::default());
        let density = 1.0;
        let velocity = Vector2::new(0.1, -0.05);

        let f_eq = solver.equilibrium_distribution(density, &velocity);

        // Calculate momentum from equilibrium distribution
        let mut momentum = Vector2::zeros();
        for i in 0..9 {
            let e_i = Vector2::new(D2Q9_VELOCITIES[i].x as f64, D2Q9_VELOCITIES[i].y as f64);
            momentum += f_eq[i] * e_i;
        }

        let expected_momentum = density * velocity;
        assert_relative_eq!(momentum.x, expected_momentum.x, epsilon = 1e-15);
        assert_relative_eq!(momentum.y, expected_momentum.y, epsilon = 1e-15);
    }
}