use criterion::{Criterion, criterion_main};
use num_valid::{
    RealScalar,
    kernels::native64_validated::{RealNative64StrictFinite, RealNative64StrictFiniteInDebug},
};
use rand::prelude::*;
use std::hint::black_box;
use try_create::TryNew;

const VECTOR_SIZE: usize = 10000;

/// Generates a vector of random f64 values for benchmarking.
fn generate_f64_data() -> Vec<f64> {
    let mut rng = rand::rng();
    (0..VECTOR_SIZE)
        .map(|_| rng.random::<f64>() * 100.0)
        .collect()
}

#[inline(always)]
fn run_summation_with_fold<T: RealScalar>(data: &[T]) -> T {
    let sum = data.iter().fold(T::zero(), |acc, x| acc + x);
    black_box(sum)
}

#[inline(always)]
fn run_summation_with_compensated_sum<T: RealScalar + std::iter::Sum>(data: &[T]) -> T {
    let sum = data.iter().cloned().sum::<T>();
    black_box(sum)
}

#[inline(always)]
fn run_function_chain<T: RealScalar>(data: &[T]) {
    let results: Vec<_> = data
        .iter()
        .cloned()
        .map(|x| x.sin().cos().abs().sqrt())
        .collect();
    black_box(results);
}

/// Benchmarks summation performance.
fn benchmark_summation(c: &mut Criterion) {
    let f64_data = generate_f64_data();

    let validated_data: Vec<_> = f64_data
        .iter()
        .map(|&x| RealNative64StrictFinite::try_new(x).unwrap())
        .collect();

    let validated_data_in_debug: Vec<_> = f64_data
        .iter()
        .map(|&x| RealNative64StrictFiniteInDebug::try_new(x).unwrap())
        .collect();

    let mut group = c.benchmark_group("Summation");

    // Benchmark 1: Raw f64 naive summation.
    group.bench_function("f64 (naive sum)", |b| {
        b.iter(|| {
            let sum = f64_data.iter().sum::<f64>();
            black_box(sum)
        })
    });

    // Benchmark 2: RealNative64StrictFinite using a simple fold.
    // This measures the overhead of the wrapper and the validated `Add` operation.
    group.bench_function("RealNative64StrictFinite (fold)", |b| {
        b.iter(|| {
            run_summation_with_fold(&validated_data);
        })
    });

    // Benchmark 3: RealNative64StrictFinite using the `Sum` trait.
    // This measures the performance of the more accurate Neumaier compensated summation.
    group.bench_function("RealNative64StrictFinite (compensated sum)", |b| {
        b.iter(|| {
            run_summation_with_compensated_sum(&validated_data);
        })
    });

    // Benchmark 4: RealNative64StrictFiniteInDebug using a simple fold.
    // This measures the overhead of the wrapper and the validated `Add` operation.
    group.bench_function("RealNative64StrictFiniteInDebug (fold)", |b| {
        b.iter(|| {
            run_summation_with_fold(&validated_data_in_debug);
        })
    });

    // Benchmark 3: RealNative64StrictFiniteInDebug using the `Sum` trait.
    // This measures the performance of the more accurate Neumaier compensated summation.
    group.bench_function("RealNative64StrictFiniteInDebug (compensated sum)", |b| {
        b.iter(|| {
            run_summation_with_compensated_sum(&validated_data_in_debug);
        })
    });

    group.finish();
}

/// Benchmarks performance of chained mathematical function calls.
fn benchmark_functions(c: &mut Criterion) {
    let f64_data = generate_f64_data();

    let validated_data: Vec<_> = f64_data
        .iter()
        .map(|&x| RealNative64StrictFinite::try_new(x.abs()).unwrap()) // Ensure positive for sqrt
        .collect();

    let validated_data_in_debug: Vec<_> = f64_data
        .iter()
        .map(|&x| RealNative64StrictFiniteInDebug::try_new(x.abs()).unwrap()) // Ensure positive for sqrt
        .collect();

    let mut group = c.benchmark_group("Function Chain (sin -> cos -> sqrt)");

    // Benchmark 1: Raw f64 function chain.
    group.bench_function("f64", |b| {
        b.iter(|| {
            run_function_chain(&f64_data);
        })
    });

    // Benchmark 2: RealNative64StrictFinite function chain.
    group.bench_function("RealNative64StrictFinite", |b| {
        b.iter(|| {
            run_function_chain(&validated_data);
        })
    });

    // Benchmark 3: RealNative64StrictFiniteInDebug function chain.
    group.bench_function("RealNative64StrictFiniteInDebug", |b| {
        b.iter(|| {
            run_function_chain(&validated_data_in_debug);
        })
    });

    group.finish();
}

#[cfg(feature = "rug")]
mod rug_benchmarks {
    use super::*;
    use num_valid::kernels::rug::RealRugStrictFinite;

    /// Generates a vector of random RealRugStrictFinite<P> values.
    fn generate_rug_data<const P: u32>() -> Vec<RealRugStrictFinite<P>> {
        let mut rng = rand::rng();
        (0..super::VECTOR_SIZE)
            .map(|_| {
                let val = rng.random::<f64>() * 100.0;
                RealRugStrictFinite::<P>::try_from_f64(val).unwrap()
            })
            .collect()
    }

    /// Generates a vector of positive random RealRugStrictFinite<P> values for functions like sqrt.
    fn generate_positive_rug_data<const P: u32>() -> Vec<RealRugStrictFinite<P>> {
        let mut rng = rand::rng();
        (0..super::VECTOR_SIZE)
            .map(|_| {
                let val = (rng.random::<f64>() * 100.0).abs();
                RealRugStrictFinite::<P>::try_from_f64(val).unwrap()
            })
            .collect()
    }

    pub fn benchmark_summation_rug(c: &mut Criterion) {
        let rug53_data = generate_rug_data::<53>();
        let rug1000_data = generate_rug_data::<1000>();

        let mut group = c.benchmark_group("Summation");

        // RealRugStrictFinite<53>
        group.bench_function("RealRugStrictFinite<53> (fold)", |b| {
            b.iter(|| {
                run_summation_with_fold(&rug53_data);
            })
        });
        group.bench_function("RealRugStrictFinite<53> (compensated sum)", |b| {
            b.iter(|| {
                run_summation_with_compensated_sum(&rug53_data);
            })
        });

        // RealRugStrictFinite<1000>
        group.bench_function("RealRugStrictFinite<1000> (fold)", |b| {
            b.iter(|| {
                run_summation_with_fold(&rug1000_data);
            })
        });
        group.bench_function("RealRugStrictFinite<1000> (compensated sum)", |b| {
            b.iter(|| {
                run_summation_with_compensated_sum(&rug1000_data);
            })
        });

        group.finish();
    }

    pub fn benchmark_functions_rug(c: &mut Criterion) {
        let rug53_data = generate_positive_rug_data::<53>();
        let rug1000_data = generate_positive_rug_data::<1000>();

        let mut group = c.benchmark_group("Function Chain (sin -> cos -> sqrt)");

        // RealRugStrictFinite<53>
        group.bench_function("RealRugStrictFinite<53>", |b| {
            b.iter(|| {
                run_function_chain(&rug53_data);
            })
        });

        // RealRugStrictFinite<1000>
        group.bench_function("RealRugStrictFinite<1000>", |b| {
            b.iter(|| {
                run_function_chain(&rug1000_data);
            })
        });

        group.finish();
    }
}

fn benches() {
    let mut c = Criterion::default().configure_from_args();
    benchmark_summation(&mut c);
    benchmark_functions(&mut c);

    #[cfg(feature = "rug")]
    {
        rug_benchmarks::benchmark_summation_rug(&mut c);
        rug_benchmarks::benchmark_functions_rug(&mut c);
    }
}

//criterion_group!(benches, benchmark_summation, benchmark_functions);
criterion_main!(benches);
