use criterion::{BatchSize, Criterion, black_box, criterion_group, criterion_main};
use num::complex::Complex;
use num::traits::Float;
use vfft::cpu::{ExecutionPlan, fft};

// Cast one float type to another, truncating if needed.
fn t0_to_t1<T0: Float, T1: Float>(val: T0) -> T1 {
    return num::cast(val).unwrap();
}

// Casts one Complex<T> to another.
fn complex_to_t<T0: Float, T1: Float>(data: &[Complex<T0>]) -> Vec<Complex<T1>> {
    data.iter()
        .map(|c| Complex::new(t0_to_t1(c.re), t0_to_t1(c.im)))
        .collect()
}

fn get_data(n: usize) -> Vec<Complex<f32>> {
    let mut data: Vec<Complex<f32>> = Vec::new();
    data.resize(n, Complex::ZERO);
    // Randomize data.
    for i in 0..data.len() {
        data[i] = Complex::new(rand::random(), rand::random());
    }
    data
}

fn do_fft_v8_benchmark_for_n(c: &mut Criterion, n: usize) {
    let data = get_data(n);
    let mut plan = ExecutionPlan::new(&data).unwrap();
    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);

    c.bench_function(format!("size {}, fft v8", n).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                fft::<f32>(result_ref, &mut plan);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_ref_benchmark_for_n(c: &mut Criterion, n: usize) {
    let data = get_data(n);

    let result_ref: Vec<rustfft::num_complex::Complex<f32>> = complex_to_t(&data);
    let mut ref_planner = rustfft::FftPlanner::<f32>::new();
    let mut scratch: Vec<rustfft::num_complex::Complex<f32>> =
        vec![Complex::<f32>::new(0.0, 0.0); n];
    let ref_fft = ref_planner.plan_fft_forward(result_ref.len());

    c.bench_function(format!("size {}, reference", n).as_str(), |b| {
        b.iter_batched_ref(
            || result_ref.clone(),
            |result_ref| {
                ref_fft.process_with_scratch(result_ref, &mut scratch);
                black_box(result_ref);
            },
            BatchSize::PerIteration,
        );
    });
}

fn do_bench(c: &mut Criterion, n: usize) {
    do_ref_benchmark_for_n(c, n);
    do_fft_v8_benchmark_for_n(c, n);
}

fn get_composites(primes: &[usize], max_number: usize) -> Vec<usize> {
    let mut composites = Vec::new();
    composites.push(1);

    for prime in primes {
        let num_before = composites.len();
        for i in 0..num_before {
            let mut composite = composites[i] * prime;
            while composite <= max_number {
                composites.push(composite);
                composite *= prime;
            }
        }
    }

    composites.retain(|&x| x >= 11);
    composites.sort();
    composites
}

pub fn criterion_benchmark(c: &mut Criterion) {
    let radices = [2, 3, 5, 7];
    const MAX_SIZE: usize = 8192;
    let composites = get_composites(&radices, MAX_SIZE);
    for composite in &composites {
        do_bench(c, *composite);
    }
}

criterion_group!(benches, criterion_benchmark);
criterion_main!(benches);