use criterion::{BatchSize, BenchmarkId, Criterion, Throughput, criterion_group, criterion_main};
use obikseq::packed_seq::PackedSeq;
use obikseq::superkmer::SuperKmer;
use obikseq::Sequence;

const LENGTHS: &[usize] = &[1, 4, 8, 16, 40, 64, 128, 255, 256];

fn make_ascii(len: usize) -> Vec<u8> {
    (0..len).map(|i| b"ACGT"[i % 4]).collect()
}

fn bench_from_ascii(c: &mut Criterion) {
    let mut group = c.benchmark_group("from_ascii");
    for &len in LENGTHS {
        let ascii = make_ascii(len);
        group.throughput(Throughput::Bytes(len as u64));
        group.bench_with_input(BenchmarkId::from_parameter(len), &ascii, |b, ascii| {
            b.iter(|| SuperKmer::from_ascii(std::hint::black_box(ascii)));
        });
    }
    group.finish();
}

fn bench_to_ascii(c: &mut Criterion) {
    let mut group = c.benchmark_group("to_ascii");
    for &len in LENGTHS {
        let sk = SuperKmer::from_ascii(&make_ascii(len));
        group.throughput(Throughput::Bytes(len as u64));
        group.bench_with_input(BenchmarkId::from_parameter(len), &sk, |b, sk| {
            b.iter(|| std::hint::black_box(sk).to_ascii());
        });
    }
    group.finish();
}

fn bench_write_ascii(c: &mut Criterion) {
    let mut group = c.benchmark_group("write_ascii");
    for &len in LENGTHS {
        let sk = SuperKmer::from_ascii(&make_ascii(len));
        group.throughput(Throughput::Bytes(len as u64));
        group.bench_with_input(BenchmarkId::from_parameter(len), &sk, |b, sk| {
            let mut buf = Vec::with_capacity(len);
            b.iter(|| {
                buf.clear();
                std::hint::black_box(sk).write_ascii(&mut buf).unwrap();
            });
        });
    }
    group.finish();
}

fn bench_revcomp(c: &mut Criterion) {
    let mut group = c.benchmark_group("revcomp");
    for &len in LENGTHS {
        let seq = PackedSeq::from_ascii(&make_ascii(len));
        group.throughput(Throughput::Bytes(len as u64));
        group.bench_with_input(BenchmarkId::from_parameter(len), &seq, |b, seq| {
            b.iter_batched(
                || seq.clone(),
                |s| { let mut s = std::hint::black_box(s); s.revcomp_inplace(); s },
                BatchSize::SmallInput,
            );
        });
    }
    group.finish();
}

fn bench_canonical(c: &mut Criterion) {
    let mut group = c.benchmark_group("canonical");
    for &len in LENGTHS {
        let sk_rc  = SuperKmer::from_ascii(&vec![b'T'; len]);
        let sk_fwd = SuperKmer::from_ascii(&vec![b'A'; len]);
        group.throughput(Throughput::Bytes(len as u64));
        group.bench_with_input(
            BenchmarkId::new("needs_revcomp", len), &sk_rc, |b, sk| {
                b.iter_batched(
                    || sk.clone(),
                    |s| std::hint::black_box(s).canonical(),
                    BatchSize::SmallInput,
                );
            }
        );
        group.bench_with_input(
            BenchmarkId::new("already_canonical", len), &sk_fwd, |b, sk| {
                b.iter_batched(
                    || sk.clone(),
                    |s| std::hint::black_box(s).canonical(),
                    BatchSize::SmallInput,
                );
            }
        );
    }
    group.finish();
}

criterion_group!(benches, bench_from_ascii, bench_to_ascii, bench_write_ascii, bench_revcomp, bench_canonical);
criterion_main!(benches);
