use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use grid1d::{
    Grid1D, IntervalPartition,
    intervals::*,
    scalars::{IntervalId, NumIntervals, PositiveNumPoints1D},
};
use num_valid::{RealNative64StrictFinite, RealNative64StrictFiniteInDebug};
use sorted_vec::partial::SortedSet;
use std::hint::black_box;
use try_create::TryNew;

#[cfg(feature = "rug")]
use num_valid::{RealRugStrictFinite, RealScalar};

// Type aliases for different scalar types
type OptimalReal = RealNative64StrictFiniteInDebug;
type SafeReal = RealNative64StrictFinite;

#[cfg(feature = "rug")]
type HighPrecisionReal64 = RealRugStrictFinite<64>;
#[cfg(feature = "rug")]
type HighPrecisionReal128 = RealRugStrictFinite<128>;
#[cfg(feature = "rug")]
type HighPrecisionReal256 = RealRugStrictFinite<256>;

/// Benchmark grid construction for different sizes and scalar types
fn bench_grid_construction(c: &mut Criterion) {
    let mut group = c.benchmark_group("grid_construction");

    // Test different grid sizes
    let sizes = [10, 100, 1_000, 10_000, 100_000];

    for &size in &sizes {
        group.throughput(Throughput::Elements(size as u64));

        // Benchmark uniform grid construction with raw f64
        group.bench_with_input(BenchmarkId::new("uniform_f64", size), &size, |b, &size| {
            let domain = IntervalClosed::new(0.0_f64, 1.0_f64);
            let num_intervals = NumIntervals::try_new(size).unwrap();

            b.iter(|| {
                let grid = Grid1D::uniform(black_box(domain.clone()), black_box(num_intervals));
                black_box(grid)
            });
        });

        // Benchmark uniform grid construction with validated scalars (debug-only validation)
        group.bench_with_input(
            BenchmarkId::new("uniform_optimal", size),
            &size,
            |b, &size| {
                let domain = IntervalClosed::new(
                    OptimalReal::try_new(0.0).unwrap(),
                    OptimalReal::try_new(1.0).unwrap(),
                );
                let num_intervals = NumIntervals::try_new(size).unwrap();

                b.iter(|| {
                    let grid = Grid1D::uniform(black_box(domain.clone()), black_box(num_intervals));
                    black_box(grid)
                });
            },
        );

        // Benchmark uniform grid construction with always-validated scalars
        group.bench_with_input(BenchmarkId::new("uniform_safe", size), &size, |b, &size| {
            let domain = IntervalClosed::new(
                SafeReal::try_new(0.0).unwrap(),
                SafeReal::try_new(1.0).unwrap(),
            );
            let num_intervals = NumIntervals::try_new(size).unwrap();

            b.iter(|| {
                let grid = Grid1D::uniform(black_box(domain.clone()), black_box(num_intervals));
                black_box(grid)
            });
        });

        #[cfg(feature = "rug")]
        {
            // Benchmark with 64-bit arbitrary precision
            group.bench_with_input(
                BenchmarkId::new("uniform_rug_64", size),
                &size,
                |b, &size| {
                    let domain = IntervalClosed::new(
                        HighPrecisionReal64::try_from_f64(0.0).unwrap(),
                        HighPrecisionReal64::try_from_f64(1.0).unwrap(),
                    );
                    let num_intervals = NumIntervals::try_new(size).unwrap();

                    b.iter(|| {
                        let grid =
                            Grid1D::uniform(black_box(domain.clone()), black_box(num_intervals));
                        black_box(grid)
                    });
                },
            );

            // Benchmark with 128-bit arbitrary precision
            group.bench_with_input(
                BenchmarkId::new("uniform_rug_128", size),
                &size,
                |b, &size| {
                    let domain = IntervalClosed::new(
                        HighPrecisionReal128::try_from_f64(0.0).unwrap(),
                        HighPrecisionReal128::try_from_f64(1.0).unwrap(),
                    );
                    let num_intervals = NumIntervals::try_new(size).unwrap();

                    b.iter(|| {
                        let grid =
                            Grid1D::uniform(black_box(domain.clone()), black_box(num_intervals));
                        black_box(grid)
                    });
                },
            );

            // Benchmark with 256-bit arbitrary precision
            group.bench_with_input(
                BenchmarkId::new("uniform_rug_256", size),
                &size,
                |b, &size| {
                    let domain = IntervalClosed::new(
                        HighPrecisionReal256::try_from_f64(0.0).unwrap(),
                        HighPrecisionReal256::try_from_f64(1.0).unwrap(),
                    );
                    let num_intervals = NumIntervals::try_new(size).unwrap();

                    b.iter(|| {
                        let grid =
                            Grid1D::uniform(black_box(domain.clone()), black_box(num_intervals));
                        black_box(grid)
                    });
                },
            );
        }

        // Benchmark non-uniform grid construction
        group.bench_with_input(
            BenchmarkId::new("non_uniform_f64", size),
            &size,
            |b, &size| {
                // Pre-generate coordinates
                let coords: Vec<f64> = (0..=size).map(|i| (i as f64) / (size as f64)).collect();

                b.iter(|| {
                    let sorted_coords = SortedSet::from_unsorted(black_box(coords.clone()));
                    let grid =
                        Grid1D::<IntervalClosed<f64>>::try_from_sorted(black_box(sorted_coords))
                            .unwrap();
                    black_box(grid)
                });
            },
        );

        #[cfg(feature = "rug")]
        {
            // Non-uniform with arbitrary precision
            group.bench_with_input(
                BenchmarkId::new("non_uniform_rug_128", size),
                &size,
                |b, &size| {
                    let coords: Vec<HighPrecisionReal128> = (0..=size)
                        .map(|i| {
                            HighPrecisionReal128::try_from_f64((i as f64) / (size as f64)).unwrap()
                        })
                        .collect();

                    b.iter(|| {
                        let sorted_coords = SortedSet::from_unsorted(black_box(coords.clone()));
                        let grid = Grid1D::<IntervalClosed<HighPrecisionReal128>>::try_from_sorted(
                            black_box(sorted_coords),
                        )
                        .unwrap();
                        black_box(grid)
                    });
                },
            );
        }
    }

    group.finish();
}

/// Benchmark point location operations - the most performance-critical operation
fn bench_point_location(c: &mut Criterion) {
    let mut group = c.benchmark_group("point_location");

    let sizes = [100, 1_000, 10_000, 100_000];
    let num_queries = 1000;

    for &size in &sizes {
        group.throughput(Throughput::Elements(num_queries));

        // Setup test data for f64
        let uniform_grid_f64 = Grid1D::uniform(
            IntervalClosed::new(0.0_f64, 1.0_f64),
            NumIntervals::try_new(size).unwrap(),
        );

        let non_uniform_coords: Vec<f64> = (0..=size)
            .map(|i| {
                let t = (i as f64) / (size as f64);
                t * t // Quadratic spacing for more realistic non-uniform distribution
            })
            .collect();
        let non_uniform_grid_f64 = Grid1D::<IntervalClosed<f64>>::try_from_sorted(
            SortedSet::from_unsorted(non_uniform_coords),
        )
        .unwrap();

        // Generate test points
        let test_points_f64: Vec<f64> = (0..num_queries)
            .map(|i| (i as f64) / (num_queries as f64))
            .collect();

        // Benchmark uniform grid point location (should be O(1))
        group.bench_with_input(
            BenchmarkId::new("uniform_f64_single", size),
            &size,
            |b, _| {
                b.iter(|| {
                    for &point in &test_points_f64 {
                        let interval_id =
                            uniform_grid_f64.find_interval_id_of_point(black_box(&point));
                        black_box(interval_id);
                    }
                });
            },
        );

        // Benchmark non-uniform grid point location (should be O(log n))
        group.bench_with_input(
            BenchmarkId::new("non_uniform_f64_single", size),
            &size,
            |b, _| {
                b.iter(|| {
                    for &point in &test_points_f64 {
                        let interval_id =
                            non_uniform_grid_f64.find_interval_id_of_point(black_box(&point));
                        black_box(interval_id);
                    }
                });
            },
        );

        // Setup test data for validated scalars
        let uniform_grid_optimal = Grid1D::uniform(
            IntervalClosed::new(
                OptimalReal::try_new(0.0).unwrap(),
                OptimalReal::try_new(1.0).unwrap(),
            ),
            NumIntervals::try_new(size).unwrap(),
        );

        let test_points_optimal: Vec<OptimalReal> = (0..num_queries)
            .map(|i| OptimalReal::try_new((i as f64) / (num_queries as f64)).unwrap())
            .collect();

        group.bench_with_input(
            BenchmarkId::new("uniform_optimal_single", size),
            &size,
            |b, _| {
                b.iter(|| {
                    for point in &test_points_optimal {
                        let interval_id =
                            uniform_grid_optimal.find_interval_id_of_point(black_box(point));
                        black_box(interval_id);
                    }
                });
            },
        );

        #[cfg(feature = "rug")]
        {
            // Setup test data for arbitrary precision
            let uniform_grid_rug = Grid1D::uniform(
                IntervalClosed::new(
                    HighPrecisionReal128::try_from_f64(0.0).unwrap(),
                    HighPrecisionReal128::try_from_f64(1.0).unwrap(),
                ),
                NumIntervals::try_new(size).unwrap(),
            );

            let test_points_rug: Vec<HighPrecisionReal128> = (0..num_queries)
                .map(|i| {
                    HighPrecisionReal128::try_from_f64((i as f64) / (num_queries as f64)).unwrap()
                })
                .collect();

            group.bench_with_input(
                BenchmarkId::new("uniform_rug_128_single", size),
                &size,
                |b, _| {
                    b.iter(|| {
                        for point in &test_points_rug {
                            let interval_id =
                                uniform_grid_rug.find_interval_id_of_point(black_box(point));
                            black_box(interval_id);
                        }
                    });
                },
            );

            // High precision 256-bit
            let uniform_grid_rug_256 = Grid1D::uniform(
                IntervalClosed::new(
                    HighPrecisionReal256::try_from_f64(0.0).unwrap(),
                    HighPrecisionReal256::try_from_f64(1.0).unwrap(),
                ),
                NumIntervals::try_new(size).unwrap(),
            );

            let test_points_rug_256: Vec<HighPrecisionReal256> = (0..num_queries)
                .map(|i| {
                    HighPrecisionReal256::try_from_f64((i as f64) / (num_queries as f64)).unwrap()
                })
                .collect();

            group.bench_with_input(
                BenchmarkId::new("uniform_rug_256_single", size),
                &size,
                |b, _| {
                    b.iter(|| {
                        for point in &test_points_rug_256 {
                            let interval_id =
                                uniform_grid_rug_256.find_interval_id_of_point(black_box(point));
                            black_box(interval_id);
                        }
                    });
                },
            );
        }

        // Benchmark batch point location
        group.bench_with_input(
            BenchmarkId::new("uniform_f64_batch", size),
            &size,
            |b, _| {
                b.iter(|| {
                    let results =
                        uniform_grid_f64.find_intervals_for_points(black_box(&test_points_f64));
                    black_box(results);
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("non_uniform_f64_batch", size),
            &size,
            |b, _| {
                b.iter(|| {
                    let results =
                        non_uniform_grid_f64.find_intervals_for_points(black_box(&test_points_f64));
                    black_box(results);
                });
            },
        );

        // Benchmark parallel batch point location (only for larger sizes)
        if size >= 1000 {
            group.bench_with_input(
                BenchmarkId::new("uniform_f64_batch_parallel", size),
                &size,
                |b, _| {
                    b.iter(|| {
                        let results = uniform_grid_f64
                            .find_intervals_for_points_parallel(black_box(&test_points_f64));
                        black_box(results);
                    });
                },
            );

            group.bench_with_input(
                BenchmarkId::new("non_uniform_f64_batch_parallel", size),
                &size,
                |b, _| {
                    b.iter(|| {
                        let results = non_uniform_grid_f64
                            .find_intervals_for_points_parallel(black_box(&test_points_f64));
                        black_box(results);
                    });
                },
            );
        }
    }

    group.finish();
}

/// Benchmark grid refinement operations
fn bench_grid_refinement(c: &mut Criterion) {
    let mut group = c.benchmark_group("grid_refinement");

    let base_sizes = [100, 1_000, 10_000];

    for &size in &base_sizes {
        group.throughput(Throughput::Elements(size as u64));

        let base_grid_f64 = Grid1D::uniform(
            IntervalClosed::new(0.0_f64, 1.0_f64),
            NumIntervals::try_new(size).unwrap(),
        );

        let base_grid_optimal = Grid1D::uniform(
            IntervalClosed::new(
                OptimalReal::try_new(0.0).unwrap(),
                OptimalReal::try_new(1.0).unwrap(),
            ),
            NumIntervals::try_new(size).unwrap(),
        );

        // Benchmark uniform refinement
        group.bench_with_input(
            BenchmarkId::new("uniform_refinement_f64", size),
            &size,
            |b, _| {
                let extra_points = PositiveNumPoints1D::try_new(1).unwrap();

                b.iter(|| {
                    let refinement = base_grid_f64
                        .clone()
                        .refine_uniform(black_box(&extra_points));
                    black_box(refinement);
                });
            },
        );

        group.bench_with_input(
            BenchmarkId::new("uniform_refinement_optimal", size),
            &size,
            |b, _| {
                let extra_points = PositiveNumPoints1D::try_new(1).unwrap();

                b.iter(|| {
                    let refinement = base_grid_optimal
                        .clone()
                        .refine_uniform(black_box(&extra_points));
                    black_box(refinement);
                });
            },
        );

        #[cfg(feature = "rug")]
        {
            let base_grid_rug = Grid1D::uniform(
                IntervalClosed::new(
                    HighPrecisionReal128::try_from_f64(0.0).unwrap(),
                    HighPrecisionReal128::try_from_f64(1.0).unwrap(),
                ),
                NumIntervals::try_new(size).unwrap(),
            );

            group.bench_with_input(
                BenchmarkId::new("uniform_refinement_rug_128", size),
                &size,
                |b, _| {
                    let extra_points = PositiveNumPoints1D::try_new(1).unwrap();

                    b.iter(|| {
                        let refinement = base_grid_rug
                            .clone()
                            .refine_uniform(black_box(&extra_points));
                        black_box(refinement);
                    });
                },
            );
        }

        // Benchmark selective refinement (refine 10% of intervals)
        group.bench_with_input(
            BenchmarkId::new("selective_refinement_f64", size),
            &size,
            |b, _| {
                let refinement_plan: std::collections::BTreeMap<_, _> = (0..size / 10)
                    .map(|i| {
                        (
                            IntervalId::new(i * 10),
                            PositiveNumPoints1D::try_new(2).unwrap(),
                        )
                    })
                    .collect();

                b.iter(|| {
                    let refinement = base_grid_f64.clone().refine(black_box(&refinement_plan));
                    black_box(refinement);
                });
            },
        );

        #[cfg(feature = "rug")]
        {
            let base_grid_rug = Grid1D::uniform(
                IntervalClosed::new(
                    HighPrecisionReal128::try_from_f64(0.0).unwrap(),
                    HighPrecisionReal128::try_from_f64(1.0).unwrap(),
                ),
                NumIntervals::try_new(size).unwrap(),
            );

            group.bench_with_input(
                BenchmarkId::new("selective_refinement_rug_128", size),
                &size,
                |b, _| {
                    let refinement_plan: std::collections::BTreeMap<_, _> = (0..size / 10)
                        .map(|i| {
                            (
                                IntervalId::new(i * 10),
                                PositiveNumPoints1D::try_new(2).unwrap(),
                            )
                        })
                        .collect();

                    b.iter(|| {
                        let refinement = base_grid_rug.clone().refine(black_box(&refinement_plan));
                        black_box(refinement);
                    });
                },
            );
        }
    }

    group.finish();
}

/// Benchmark grid union operations
fn bench_grid_union(c: &mut Criterion) {
    let mut group = c.benchmark_group("grid_union");

    let sizes = [100, 1_000, 10_000];

    for &size in &sizes {
        group.throughput(Throughput::Elements(size as u64));

        // Create two grids with different spacing patterns for f64
        let grid_a_f64 = Grid1D::uniform(
            IntervalClosed::new(0.0_f64, 1.0_f64),
            NumIntervals::try_new(size).unwrap(),
        );

        let coords_b_f64: Vec<f64> = (0..=size)
            .map(|i| {
                let t = (i as f64) / (size as f64);
                t * t.sqrt() // Different spacing pattern
            })
            .collect();
        let grid_b_f64 =
            Grid1D::<IntervalClosed<f64>>::try_from_sorted(SortedSet::from_unsorted(coords_b_f64))
                .unwrap();

        group.bench_with_input(
            BenchmarkId::new("union_creation_f64", size),
            &size,
            |b, _| {
                b.iter(|| {
                    let union = grid1d::Grid1DUnion::try_new(
                        black_box(&grid_a_f64),
                        black_box(&grid_b_f64),
                    )
                    .unwrap();
                    black_box(union);
                });
            },
        );

        // Optimal real type
        let grid_a_optimal = Grid1D::uniform(
            IntervalClosed::new(
                OptimalReal::try_new(0.0).unwrap(),
                OptimalReal::try_new(1.0).unwrap(),
            ),
            NumIntervals::try_new(size).unwrap(),
        );

        let coords_b_optimal: Vec<OptimalReal> = (0..=size)
            .map(|i| {
                let t = (i as f64) / (size as f64);
                OptimalReal::try_new(t * t.sqrt()).unwrap() // Different spacing pattern
            })
            .collect();
        let grid_b_optimal = Grid1D::<IntervalClosed<OptimalReal>>::try_from_sorted(
            SortedSet::from_unsorted(coords_b_optimal),
        )
        .unwrap();

        group.bench_with_input(
            BenchmarkId::new("union_creation_optimal", size),
            &size,
            |b, _| {
                b.iter(|| {
                    let union = grid1d::Grid1DUnion::try_new(
                        black_box(&grid_a_optimal),
                        black_box(&grid_b_optimal),
                    )
                    .unwrap();
                    black_box(union);
                });
            },
        );

        #[cfg(feature = "rug")]
        {
            let grid_a_rug = Grid1D::uniform(
                IntervalClosed::new(
                    HighPrecisionReal128::try_from_f64(0.0).unwrap(),
                    HighPrecisionReal128::try_from_f64(1.0).unwrap(),
                ),
                NumIntervals::try_new(size).unwrap(),
            );

            let coords_b_rug: Vec<HighPrecisionReal128> = (0..=size)
                .map(|i| {
                    let t = (i as f64) / (size as f64);
                    HighPrecisionReal128::try_from_f64(t * t.sqrt()).unwrap() // Different spacing pattern
                })
                .collect();
            let grid_b_rug = Grid1D::<IntervalClosed<HighPrecisionReal128>>::try_from_sorted(
                SortedSet::from_unsorted(coords_b_rug),
            )
            .unwrap();

            group.bench_with_input(
                BenchmarkId::new("union_creation_rug_128", size),
                &size,
                |b, _| {
                    b.iter(|| {
                        let union = grid1d::Grid1DUnion::try_new(
                            black_box(&grid_a_rug),
                            black_box(&grid_b_rug),
                        )
                        .unwrap();
                        black_box(union);
                    });
                },
            );
        }
    }

    group.finish();
}

// Group all benchmarks
criterion_group!(
    benches,
    bench_grid_construction,
    bench_point_location,
    bench_grid_refinement,
    bench_grid_union
);
criterion_main!(benches);
