//! Performance benchmarks for RTX CFD solvers
//!
//! Benchmarks key computational kernels and full solver performance
//! across different problem sizes and configurations.

use criterion::{BenchmarkId, Criterion, black_box, criterion_group, criterion_main};
use nalgebra::DVector;
use rtx_cfd::{
    discretization::{
        DifferencingScheme, FiniteDifferenceMethod, FiniteVolumeMethod, FluxScheme, GridSpacing,
        SpatialOrder,
    },
    solvers::incompressible::{FlowField, SimpleAlgorithm},
    turbulence::{KEpsilonModel, KEpsilonVariant, SmagorinskyModel, TurbulenceState},
};

/// Benchmark finite volume method discretization
fn bench_fvm_discretization(c: &mut Criterion) {
    let mut group = c.benchmark_group("FVM Discretization");

    for size in [100, 500, 1000, 2000].iter() {
        group.benchmark_with_input(BenchmarkId::new("scalar", size), size, |b, &size| {
            let mut fvm = FiniteVolumeMethod::new(SpatialOrder::Second, FluxScheme::Central)
                .with_diffusion_coefficient(1e-3);

            // Create 1D mesh
            for i in 0..size {
                fvm.add_cell(1.0, [i as f64, 0.0, 0.0]);
            }

            for i in 0..size - 1 {
                fvm.add_face([1.0, 0.0, 0.0], 1.0, i, Some(i + 1)).unwrap();
            }

            let phi = DVector::zeros(size);
            let velocity = DVector::zeros(size - 1);

            b.iter(|| {
                let matrix = fvm.discretize_scalar(black_box(&phi), black_box(&velocity));
                black_box(matrix)
            });
        });

        group.benchmark_with_input(BenchmarkId::new("momentum", size), size, |b, &size| {
            let mut fvm = FiniteVolumeMethod::new(SpatialOrder::Second, FluxScheme::Upwind)
                .with_diffusion_coefficient(1e-3);

            for i in 0..size {
                fvm.add_cell(1.0, [i as f64, 0.0, 0.0]);
            }

            for i in 0..size - 1 {
                fvm.add_face([1.0, 0.0, 0.0], 1.0, i, Some(i + 1)).unwrap();
            }

            let velocity = DVector::ones(size);
            let pressure = DVector::zeros(size);

            b.iter(|| {
                let matrix = fvm.discretize_momentum(black_box(&velocity), black_box(&pressure));
                black_box(matrix)
            });
        });

        group.benchmark_with_input(
            BenchmarkId::new("flux_calculation", size),
            size,
            |b, &size| {
                let mut fvm = FiniteVolumeMethod::new(SpatialOrder::Second, FluxScheme::Central);

                for i in 0..size {
                    fvm.add_cell(1.0, [i as f64, 0.0, 0.0]);
                }

                for i in 0..size - 1 {
                    fvm.add_face([1.0, 0.0, 0.0], 1.0, i, Some(i + 1)).unwrap();
                }

                let phi = DVector::from_fn(size, |i, _| (i as f64).sin());
                let velocity = DVector::ones(size - 1);

                b.iter(|| {
                    let fluxes = fvm.calculate_fluxes(black_box(&phi), black_box(&velocity));
                    black_box(fluxes)
                });
            },
        );
    }

    group.finish();
}

/// Benchmark finite difference method
fn bench_fdm_discretization(c: &mut Criterion) {
    let mut group = c.benchmark_group("FDM Discretization");

    for grid_size in [32, 64, 128, 256].iter() {
        let total_size = grid_size * grid_size;

        group.benchmark_with_input(
            BenchmarkId::new("laplacian_2d", grid_size),
            &total_size,
            |b, &size| {
                let spacing = GridSpacing::uniform(1.0 / *grid_size as f64);
                let fdm = FiniteDifferenceMethod::new(
                    SpatialOrder::Second,
                    DifferencingScheme::Central,
                    spacing,
                    [*grid_size, *grid_size, 1],
                )
                .unwrap();

                b.iter(|| {
                    let matrix = fdm.build_laplacian_2d();
                    black_box(matrix)
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("derivative_1d", grid_size),
            grid_size,
            |b, &size| {
                let spacing = GridSpacing::uniform(1.0 / size as f64);
                let fdm = FiniteDifferenceMethod::new(
                    SpatialOrder::Second,
                    DifferencingScheme::Central,
                    spacing,
                    [size, 1, 1],
                )
                .unwrap();

                b.iter(|| {
                    let matrix = fdm.build_derivative_matrix_1d(size, 1.0 / size as f64, 1);
                    black_box(matrix)
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("scalar_discretization", grid_size),
            &total_size,
            |b, &size| {
                let spacing = GridSpacing::uniform(1.0 / *grid_size as f64);
                let fdm = FiniteDifferenceMethod::new(
                    SpatialOrder::Second,
                    DifferencingScheme::Central,
                    spacing,
                    [*grid_size, *grid_size, 1],
                )
                .unwrap();

                let phi = DVector::zeros(size);
                let velocity = DVector::zeros(size);

                b.iter(|| {
                    let matrix = fdm.discretize_scalar(black_box(&phi), black_box(&velocity));
                    black_box(matrix)
                });
            },
        );
    }

    group.finish();
}

/// Benchmark turbulence models
fn bench_turbulence_models(c: &mut Criterion) {
    let mut group = c.benchmark_group("Turbulence Models");

    for size in [500, 1000, 2000, 5000].iter() {
        group.benchmark_with_input(
            BenchmarkId::new("k_epsilon_standard", size),
            size,
            |b, &size| {
                let mut model = KEpsilonModel::new(KEpsilonVariant::Standard, size);
                let mut state = TurbulenceState::new(size);
                state.initialize_k_epsilon(1e-6, 1e-8);

                model.initialize_from_state(&state).unwrap();

                // Add some velocity gradients for realistic computation
                for i in 0..size.min(state.velocity_gradients.len()) {
                    state.velocity_gradients[i][0][1] = 1.0; // du/dy
                }

                b.iter(|| {
                    model.update(black_box(&state), black_box(1e-3)).unwrap();
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("k_epsilon_realizable", size),
            size,
            |b, &size| {
                let mut model = KEpsilonModel::new(KEpsilonVariant::Realizable, size);
                let mut state = TurbulenceState::new(size);
                state.initialize_k_epsilon(1e-6, 1e-8);

                model.initialize_from_state(&state).unwrap();

                for i in 0..size.min(state.velocity_gradients.len()) {
                    state.velocity_gradients[i][0][1] = 1.0;
                }

                b.iter(|| {
                    model.update(black_box(&state), black_box(1e-3)).unwrap();
                });
            },
        );

        group.benchmark_with_input(BenchmarkId::new("smagorinsky", size), size, |b, &size| {
            let mut model = SmagorinskyModel::new(size);
            let filter_width = DVector::from_element(size, 0.1);
            model.set_filter_width(filter_width).unwrap();

            let mut state = TurbulenceState::new(size);
            for i in 0..size.min(state.velocity_gradients.len()) {
                state.velocity_gradients[i][0][1] = 1.0;
            }

            b.iter(|| {
                model.update(black_box(&state), black_box(1e-3)).unwrap();
            });
        });

        group.benchmark_with_input(
            BenchmarkId::new("production_terms", size),
            size,
            |b, &size| {
                let model = KEpsilonModel::new(KEpsilonVariant::Standard, size);
                let mut state = TurbulenceState::new(size);
                state.initialize_k_epsilon(1e-6, 1e-8);

                for i in 0..size.min(state.velocity_gradients.len()) {
                    state.velocity_gradients[i][0][1] = 1.0;
                }

                b.iter(|| {
                    let production = model.production_terms(black_box(&state));
                    black_box(production)
                });
            },
        );
    }

    group.finish();
}

/// Benchmark flow field operations
fn bench_flow_field_operations(c: &mut Criterion) {
    let mut group = c.benchmark_group("Flow Field Operations");

    for size in [1000, 5000, 10000, 20000].iter() {
        group.benchmark_with_input(
            BenchmarkId::new("velocity_operations", size),
            size,
            |b, &size| {
                let mut flow_field = FlowField::new(size);

                // Initialize with some values
                for i in 0..size {
                    flow_field
                        .set_velocity(i, [i as f64, (i * 2) as f64, 0.0])
                        .unwrap();
                }

                b.iter(|| {
                    for i in 0..size {
                        let vel = flow_field.get_velocity(black_box(i)).unwrap();
                        black_box(vel);
                    }
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("pressure_operations", size),
            size,
            |b, &size| {
                let mut flow_field = FlowField::new(size);

                for i in 0..size {
                    flow_field.set_pressure(i, i as f64 * 0.1).unwrap();
                }

                b.iter(|| {
                    for i in 0..size {
                        let pressure = flow_field.get_pressure(black_box(i)).unwrap();
                        black_box(pressure);
                    }
                });
            },
        );

        group.benchmark_with_input(BenchmarkId::new("statistics", size), size, |b, &size| {
            let mut flow_field = FlowField::new(size);

            // Initialize with realistic velocity field
            for i in 0..size {
                let u = (i as f64 / size as f64).sin();
                let v = (i as f64 / size as f64 * 2.0).cos();
                flow_field.set_velocity(i, [u, v, 0.0]).unwrap();
            }

            b.iter(|| {
                let max_vel = flow_field.calculate_max_velocity_magnitude();
                let ke = flow_field.calculate_kinetic_energy();
                black_box((max_vel, ke));
            });
        });
    }

    group.finish();
}

/// Benchmark SIMPLE algorithm
fn bench_simple_algorithm(c: &mut Criterion) {
    let mut group = c.benchmark_group("SIMPLE Algorithm");

    for grid_size in [32, 64, 128].iter() {
        let n_cells = grid_size * grid_size;

        group.benchmark_with_input(
            BenchmarkId::new("solve_step", grid_size),
            &n_cells,
            |b, &size| {
                let mut flow_field = FlowField::new(size);

                // Initialize with some non-trivial field
                for i in 0..size {
                    let u = 1.0 + 0.1 * (i as f64).sin();
                    let v = 0.1 * (i as f64 * 2.0).cos();
                    flow_field.set_velocity(i, [u, v, 0.0]).unwrap();
                    flow_field.set_pressure(i, 0.0).unwrap();
                }

                let discretization =
                    FiniteVolumeMethod::new(SpatialOrder::Second, FluxScheme::Upwind)
                        .with_diffusion_coefficient(1e-3);

                let mut solver = SimpleAlgorithm::new()
                    .with_max_iterations(10) // Limit iterations for benchmarking
                    .with_tolerance(1e-6)
                    .with_under_relaxation(0.7, 0.3);

                b.iter(|| {
                    let residual =
                        solver.solve_step(black_box(&mut flow_field), black_box(&discretization));
                    black_box(residual)
                });
            },
        );
    }

    group.finish();
}

/// Benchmark different flux schemes
fn bench_flux_schemes(c: &mut Criterion) {
    let mut group = c.benchmark_group("Flux Schemes");

    let size = 1000;
    let schemes = [
        ("central", FluxScheme::Central),
        ("upwind", FluxScheme::Upwind),
        ("quick", FluxScheme::Quick),
        ("power_law", FluxScheme::PowerLaw),
    ];

    for (name, scheme) in schemes.iter() {
        group.benchmark_with_input(
            BenchmarkId::new("flux_calculation", name),
            scheme,
            |b, &scheme| {
                let mut fvm = FiniteVolumeMethod::new(SpatialOrder::Second, scheme)
                    .with_diffusion_coefficient(1e-3);

                for i in 0..size {
                    fvm.add_cell(1.0, [i as f64, 0.0, 0.0]);
                }

                for i in 0..size - 1 {
                    fvm.add_face([1.0, 0.0, 0.0], 1.0, i, Some(i + 1)).unwrap();
                }

                let phi = DVector::from_fn(size, |i, _| (i as f64 * 0.1).sin());
                let velocity =
                    DVector::from_fn(size - 1, |i, _| 1.0 + 0.5 * (i as f64 * 0.2).cos());

                b.iter(|| {
                    let fluxes = fvm.calculate_fluxes(black_box(&phi), black_box(&velocity));
                    black_box(fluxes)
                });
            },
        );
    }

    group.finish();
}

/// Benchmark memory allocation patterns
fn bench_memory_allocation(c: &mut Criterion) {
    let mut group = c.benchmark_group("Memory Allocation");

    for size in [1000, 5000, 10000].iter() {
        group.benchmark_with_input(
            BenchmarkId::new("flow_field_creation", size),
            size,
            |b, &size| {
                b.iter(|| {
                    let flow_field = FlowField::new(black_box(size));
                    black_box(flow_field);
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("turbulence_state_creation", size),
            size,
            |b, &size| {
                b.iter(|| {
                    let mut state = TurbulenceState::new(black_box(size));
                    state.initialize_k_epsilon(1e-6, 1e-8);
                    black_box(state);
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("k_epsilon_creation", size),
            size,
            |b, &size| {
                b.iter(|| {
                    let model = KEpsilonModel::new(KEpsilonVariant::Standard, black_box(size));
                    black_box(model);
                });
            },
        );
    }

    group.finish();
}

/// Benchmark scaling with problem size
fn bench_scaling(c: &mut Criterion) {
    let mut group = c.benchmark_group("Scaling");

    // Test how performance scales with problem size
    let sizes = [100, 200, 500, 1000, 2000, 5000];

    for size in sizes.iter() {
        group.benchmark_with_input(
            BenchmarkId::new("total_fvm_workflow", size),
            size,
            |b, &size| {
                b.iter(|| {
                    // Complete FVM workflow
                    let mut fvm =
                        FiniteVolumeMethod::new(SpatialOrder::Second, FluxScheme::Central);

                    for i in 0..size {
                        fvm.add_cell(1.0, [i as f64, 0.0, 0.0]);
                    }

                    for i in 0..size - 1 {
                        fvm.add_face([1.0, 0.0, 0.0], 1.0, i, Some(i + 1)).unwrap();
                    }

                    let phi = DVector::from_fn(size, |i, _| (i as f64 * 0.1).sin());
                    let velocity = DVector::ones(size - 1);

                    let matrix = fvm.discretize_scalar(&phi, &velocity).unwrap();
                    let fluxes = fvm.calculate_fluxes(&phi, &velocity).unwrap();

                    black_box((matrix, fluxes));
                });
            },
        );

        group.benchmark_with_input(
            BenchmarkId::new("total_turbulence_workflow", size),
            size,
            |b, &size| {
                b.iter(|| {
                    // Complete turbulence modeling workflow
                    let mut model = KEpsilonModel::new(KEpsilonVariant::Standard, size);
                    let mut state = TurbulenceState::new(size);
                    state.initialize_k_epsilon(1e-6, 1e-8);

                    model.initialize_from_state(&state).unwrap();

                    for i in 0..size.min(state.velocity_gradients.len()) {
                        state.velocity_gradients[i][0][1] = 1.0;
                    }

                    let production = model.production_terms(&state).unwrap();
                    model.update(&state, 1e-3).unwrap();
                    let nu_t = model.turbulent_viscosity(&state).unwrap();

                    black_box((production, nu_t));
                });
            },
        );
    }

    group.finish();
}

criterion_group!(
    benches,
    bench_fvm_discretization,
    bench_fdm_discretization,
    bench_turbulence_models,
    bench_flow_field_operations,
    bench_simple_algorithm,
    bench_flux_schemes,
    bench_memory_allocation,
    bench_scaling
);

criterion_main!(benches);