//! Tests for LBM boundary conditions

use approx::assert_relative_eq;
use nalgebra::Vector2;
use rtx_cfd::solvers::lbm::boundary::{
    BounceBackBc, BoundaryOrientation, LbmBoundaryCondition, ZouHeBc,
};
use rtx_cfd::solvers::lbm::{D2Q9Parameters, D2Q9Solver};

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

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

    let initial_mass = solver.total_mass();

    // Run simulation with bounce-back boundaries
    for _ in 0..100 {
        solver.step(); // Uses bounce-back internally
    }

    let final_mass = solver.total_mass();

    // Mass should be conserved with bounce-back
    assert_relative_eq!(final_mass, initial_mass, epsilon = 1e-10);
}

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

    // Initialize with uniform flow
    solver.initialize_uniform(1.0, Vector2::new(0.1, 0.0));

    // Run for sufficient time to establish boundary layer
    for _ in 0..500 {
        solver.step();
    }

    // Check that velocity is approximately zero at walls
    for x in 0..nx {
        // Bottom wall
        let vars_bottom = solver.macroscopic_variables_at(x, 0);
        assert!(
            vars_bottom.velocity.norm() < 1e-3,
            "Bottom wall velocity should be small: {}",
            vars_bottom.velocity.norm()
        );

        // Top wall
        let vars_top = solver.macroscopic_variables_at(x, ny - 1);
        assert!(
            vars_top.velocity.norm() < 1e-3,
            "Top wall velocity should be small: {}",
            vars_top.velocity.norm()
        );
    }
}

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

    // Initialize equilibrium state
    solver.initialize_uniform(1.0, Vector2::zeros());

    // Create specific distribution functions for testing
    let mut f = vec![0.0; 9];
    f[0] = 4.0 / 9.0; // Rest particle
    f[1] = 1.0 / 18.0; // East velocity (should bounce back to west)
    f[2] = 1.0 / 18.0; // North velocity (should bounce back)
    f[3] = 1.0 / 18.0; // West velocity
    f[4] = 1.0 / 18.0; // South velocity
    f[5] = 1.0 / 72.0; // Northeast velocity
    f[6] = 1.0 / 72.0; // Northwest velocity
    f[7] = 1.0 / 72.0; // Southwest velocity
    f[8] = 1.0 / 72.0; // Southeast velocity

    // Set this distribution at a wall cell
    solver.set_distribution_at(0, 2, &f);

    // Store original values
    let f_original = solver.distribution_at(0, 2);

    // Apply bounce-back boundary conditions
    solver.apply_bounce_back_boundaries();

    let f_after = solver.distribution_at(0, 2);

    // Check bounce-back behavior for wall boundaries
    // Bottom wall: north-facing velocities should bounce back
    // The specific behavior depends on which wall we're testing
    assert!(f_after[0] == f_original[0]); // Rest particle unchanged
}

#[test]
fn test_zou_he_velocity_boundary() {
    let bc = ZouHeBc::new(0.1); // Velocity boundary

    // Create test distribution
    let mut f = vec![1.0 / 9.0; 9]; // Start with uniform distribution

    // Apply Zou-He boundary condition
    bc.apply(&mut f, 0, 0);

    // Check that the distribution is modified
    // (The exact check depends on implementation details)
    assert!(f.iter().sum::<f64>() > 0.0);
}

#[test]
fn test_zou_he_pressure_boundary() {
    // A pressure boundary prescribes the density, so it must be built with
    // the pressure constructor.
    //
    // This previously used `ZouHeBc::new(1.2)`, which is the legacy
    // constructor for a *velocity* boundary. That imposed a lattice velocity
    // of 1.2 — well above the lattice speed of sound, 1/sqrt(3) — and the
    // Zou-He density relation `rho = (...) / (1 - u)` then divides by -0.2 and
    // returns a negative density. The distribution came back negative, which
    // is what the "sum > 0" assertion was detecting.
    let target_density = 1.2;
    let bc = ZouHeBc::pressure(target_density, BoundaryOrientation::Left);

    let mut f = vec![1.0 / 9.0; 9];
    bc.apply(&mut f, 0, 0);

    // The contract of a pressure boundary is the density it produces, so
    // assert that rather than merely that the result is positive.
    assert_relative_eq!(f.iter().sum::<f64>(), target_density, epsilon = 1e-12);
}

#[test]
fn test_boundary_condition_trait() {
    let bounce_back = BounceBackBc;
    let zou_he = ZouHeBc::new(0.05);

    let mut f1 = vec![1.0 / 9.0; 9];
    let mut f2 = f1.clone();

    // Both should implement the trait
    bounce_back.apply(&mut f1, 0, 0);
    zou_he.apply(&mut f2, 0, 0);

    // Results should be different
    assert_ne!(f1, f2);
}

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

    // Initialize with zero velocity
    solver.initialize_uniform(1.0, Vector2::zeros());

    // Run a few steps to establish flow
    for _ in 0..10 {
        solver.step_with_boundaries(|s| {
            // Apply no-slip on walls
            s.apply_bounce_back_boundaries();

            // Apply moving lid condition (simplified - just set velocity)
            let lid_velocity = 0.1;
            for x in 1..(nx - 1) {
                let density = 1.0;
                let velocity = Vector2::new(lid_velocity, 0.0);
                let f_eq = s.equilibrium_distribution(density, &velocity);
                s.set_distribution_at(x, ny - 1, &f_eq);
            }
        });
    }

    // Populations move exactly one lattice cell per step, so after 10 steps
    // the lid can only have influenced the 10 rows beneath it. Flow must have
    // developed there.
    let steps = 10;
    let near_lid = solver.macroscopic_variables_at(nx / 2, ny - 2);
    assert!(
        near_lid.velocity.norm() > 0.0,
        "no flow beneath the lid after {steps} steps"
    );

    // ...and must not have appeared beyond the front. The domain centre is
    // 16 rows from the lid, so it is still exactly at rest.
    //
    // The assertion here previously demanded non-zero velocity at the centre
    // after 10 steps, which would require information to travel 16 cells in
    // 10 -- faster than the lattice permits. Asserting the propagation front
    // in both directions tests something the solver could actually get wrong.
    let centre = solver.macroscopic_variables_at(nx / 2, ny / 2);
    assert!(
        centre.velocity.norm() < 1e-12,
        "the lid influenced the domain centre {} rows away in only {steps} steps, \
         which exceeds the lattice propagation speed: {}",
        ny - 1 - ny / 2,
        centre.velocity.norm()
    );
}