//! Tests for D2Q9 Lattice Boltzmann Method implementation

use approx::assert_relative_eq;
use nalgebra::Vector2;
use rtx_cfd::solvers::lbm::common::MacroscopicVariables;
use rtx_cfd::solvers::lbm::{D2Q9Parameters, D2Q9Solver};

#[test]
fn test_d2q9_lattice_velocities() {
    let solver = D2Q9Solver::new(10, 10, D2Q9Parameters::default());
    let velocities = solver.lattice_velocities();

    // D2Q9 should have 9 velocities
    assert_eq!(velocities.len(), 9);

    // Check specific velocity directions
    assert_eq!(velocities[0], Vector2::new(0, 0)); // Rest
    assert_eq!(velocities[1], Vector2::new(1, 0)); // East
    assert_eq!(velocities[2], Vector2::new(0, 1)); // North
    assert_eq!(velocities[3], Vector2::new(-1, 0)); // West
    assert_eq!(velocities[4], Vector2::new(0, -1)); // South
    assert_eq!(velocities[5], Vector2::new(1, 1)); // Northeast
    assert_eq!(velocities[6], Vector2::new(-1, 1)); // Northwest
    assert_eq!(velocities[7], Vector2::new(-1, -1)); // Southwest
    assert_eq!(velocities[8], Vector2::new(1, -1)); // Southeast
}

#[test]
fn test_d2q9_weights() {
    let solver = D2Q9Solver::new(10, 10, D2Q9Parameters::default());
    let weights = solver.weights();

    // D2Q9 should have 9 weights
    assert_eq!(weights.len(), 9);

    // Check weight values
    assert_relative_eq!(weights[0], 4.0 / 9.0, epsilon = 1e-12); // Rest particle
    assert_relative_eq!(weights[1], 1.0 / 9.0, epsilon = 1e-12); // Cardinal directions
    assert_relative_eq!(weights[2], 1.0 / 9.0, epsilon = 1e-12);
    assert_relative_eq!(weights[3], 1.0 / 9.0, epsilon = 1e-12);
    assert_relative_eq!(weights[4], 1.0 / 9.0, epsilon = 1e-12);
    assert_relative_eq!(weights[5], 1.0 / 36.0, epsilon = 1e-12); // Diagonal directions
    assert_relative_eq!(weights[6], 1.0 / 36.0, epsilon = 1e-12);
    assert_relative_eq!(weights[7], 1.0 / 36.0, epsilon = 1e-12);
    assert_relative_eq!(weights[8], 1.0 / 36.0, epsilon = 1e-12);

    // Weights should sum to 1
    let sum: f64 = weights.iter().sum();
    assert_relative_eq!(sum, 1.0, epsilon = 1e-12);
}

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

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

    // Should have 9 components
    assert_eq!(f_eq.len(), 9);

    // All components should be positive
    for &val in &f_eq {
        assert!(
            val > 0.0,
            "Equilibrium distribution component should be positive"
        );
    }

    // Sum should equal density
    let sum: f64 = f_eq.iter().sum();
    assert_relative_eq!(sum, density, epsilon = 1e-12);
}

#[test]
fn test_d2q9_equilibrium_at_rest() {
    let solver = D2Q9Solver::new(10, 10, D2Q9Parameters::default());
    let density = 1.0;
    let velocity = Vector2::zeros();

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

    // For zero velocity, equilibrium should be density * weight
    for i in 0..9 {
        assert_relative_eq!(f_eq[i], density * weights[i], epsilon = 1e-12);
    }
}

#[test]
fn test_d2q9_macroscopic_variables() {
    let mut solver = D2Q9Solver::new(10, 10, D2Q9Parameters::default());

    // Initialize with known state
    let density = 1.2;
    let velocity = Vector2::new(0.1, -0.05);

    // Set distribution functions to equilibrium
    let f_eq = solver.equilibrium_distribution(density, &velocity);
    solver.set_distribution_at(5, 5, &f_eq);

    // Extract macroscopic variables
    let macro_vars = solver.macroscopic_variables_at(5, 5);

    assert_relative_eq!(macro_vars.density, density, epsilon = 1e-12);
    assert_relative_eq!(macro_vars.velocity.x, velocity.x, epsilon = 1e-12);
    assert_relative_eq!(macro_vars.velocity.y, velocity.y, epsilon = 1e-12);
}

#[test]
fn test_d2q9_bgk_collision() {
    let mut solver = D2Q9Solver::new(10, 10, D2Q9Parameters::new(0.6)); // tau = 0.6

    // Initialize with non-equilibrium state
    let density = 1.0;
    let velocity = Vector2::new(0.1, 0.0);
    let f_eq = solver.equilibrium_distribution(density, &velocity);

    // Perturb from equilibrium
    let mut f = f_eq.clone();
    f[1] += 0.1; // Add perturbation

    solver.set_distribution_at(5, 5, &f);

    // Apply collision step
    solver.collision_step();

    // Check that distribution moves toward equilibrium
    let f_after = solver.distribution_at(5, 5);

    // The perturbed component should be closer to equilibrium
    assert!(f_after[1] < f[1]);
    assert!(f_after[1] > f_eq[1]);
}

#[test]
fn test_d2q9_streaming_step() {
    let mut solver = D2Q9Solver::new(5, 5, D2Q9Parameters::default());

    // Initialize central cell with specific distribution
    let mut f = vec![0.0; 9];
    f[1] = 1.0; // Only eastward component
    solver.set_distribution_at(2, 2, &f);

    // Apply streaming step
    solver.streaming_step();

    // Check that the distribution has moved eastward
    let f_east = solver.distribution_at(3, 2);
    assert_relative_eq!(f_east[1], 1.0, epsilon = 1e-12);

    // Original cell should have zero eastward component
    let f_original = solver.distribution_at(2, 2);
    assert_relative_eq!(f_original[1], 0.0, epsilon = 1e-12);
}

#[test]
fn test_d2q9_poiseuille_flow_analytical() {
    // Test against analytical solution for Poiseuille flow
    let nx = 32;
    let ny = 32;
    let mut solver = D2Q9Solver::new(nx, ny, D2Q9Parameters::new(0.8));

    // Initialize Poiseuille flow
    solver.initialize_poiseuille_flow(0.01); // Small driving force

    // Run for sufficient time to reach steady state with equilibrium BCs
    for _ in 0..1000 {
        solver.step_with_boundaries(|s| s.apply_no_slip_boundaries());
    }

    // Check parabolic velocity profile at center
    let x_center = nx / 2;
    let mut max_velocity = 0.0f64;

    for y in 1..(ny - 1) {
        let vars = solver.macroscopic_variables_at(x_center, y);
        max_velocity = max_velocity.max(vars.velocity.x);
    }

    // For Poiseuille flow, velocity should be maximum at center
    let center_vars = solver.macroscopic_variables_at(x_center, ny / 2);
    assert_relative_eq!(center_vars.velocity.x, max_velocity, epsilon = 1e-2);

    // Velocity should be zero at walls
    let wall_bottom = solver.macroscopic_variables_at(x_center, 0);
    let wall_top = solver.macroscopic_variables_at(x_center, ny - 1);
    assert!(wall_bottom.velocity.x.abs() < 1e-6);
    assert!(wall_top.velocity.x.abs() < 1e-6);
}

#[test]
fn test_d2q9_mass_conservation() {
    let nx = 16;
    let ny = 16;
    let mut solver = D2Q9Solver::new(nx, ny, D2Q9Parameters::default());

    // Initialize with uniform density
    solver.initialize_uniform(1.0, Vector2::new(0.05, 0.02));

    let initial_mass = solver.total_mass();

    // Run simulation for several steps with periodic boundaries
    for _ in 0..100 {
        solver.step_periodic();
    }

    let final_mass = solver.total_mass();

    // Mass should be conserved
    assert_relative_eq!(final_mass, initial_mass, epsilon = 1e-12);
}

#[test]
fn test_d2q9_parameters_validation() {
    // Valid parameters
    let valid_params = D2Q9Parameters::new(0.6);
    assert!(valid_params.validate().is_ok());

    // Invalid relaxation time (too small)
    let invalid_params = D2Q9Parameters::new(0.4);
    assert!(invalid_params.validate().is_err());

    // Invalid relaxation time (too large)
    let invalid_params = D2Q9Parameters::new(2.1);
    assert!(invalid_params.validate().is_err());
}