use std::iter;

fn optimized_sieve(limit: usize) -> impl Iterator<Item = usize> {
    let map_func = |cmpsts: Vec<u32>| {
        move |i| {
            if i < 0 {
                Some(2)
            } else {
                if cmpsts[i as usize >> 5] & (1u32 << (i & 31)) == 0 {
                    Some((i + i + 3) as usize)
                } else {
                    None
                }
            }
        }
    };

    // Weird-looking exprs so types line up
    if limit == 2 {
        return (0..0).into_iter().filter_map(map_func(vec![]));
    } else if limit < 2 {
        return (-1..0).into_iter().filter_map(map_func(vec![]));
    }

    let ndxlmt = (limit - 3) / 2 + 1;
    let bfsz = ((limit - 3) / 2) / 32 + 1;
    let mut cmpsts = vec![0u32; bfsz];
    let sqrtndxlmt = ((limit as f64).sqrt() as usize - 3) / 2 + 1;

    for ndx in 0..sqrtndxlmt {
        if (cmpsts[ndx >> 5] & (1u32 << (ndx & 31))) == 0 {
            let p = ndx + ndx + 3;
            let mut cullpos = (p * p - 3) / 2;
            while cullpos < ndxlmt {
                cmpsts[cullpos >> 5] |= 1u32 << (cullpos & 31);
                cullpos += p;
            }
        }
    }

    (-1..ndxlmt as isize)
        .into_iter()
        .filter_map(map_func(cmpsts))
}

fn nth_prime_upper_bound(n: usize) -> Option<usize> {
    if let Some(upper) = [0, 2, 3, 5, 7, 11].get(n) {
        return Some(*upper);
    }

    let n_f64 = n as f64;
    let logn = n_f64.log2();
    let loglogn = logn.log2();

    let upper = match n {
        n if n >= 688383 => n_f64 * (logn + loglogn - 1. + ((loglogn - 2.) / logn)),
        n if n >= 178974 => n_f64 * (logn + loglogn - 1. + ((loglogn - 1.95) / logn)),
        n if n >= 39017 => n_f64 * (logn + loglogn - 1. + ((loglogn - 0.9484) / logn)),
        _ => n_f64 * (logn + 0.6 * loglogn),
    };

    let upper = upper.ceil();

    if upper <= std::usize::MAX as f64 {
        Some(upper as _)
    } else {
        None
    }
}

#[no_mangle]
pub extern "C" fn nth_prime(n: usize) -> usize {
    optimized_sieve(nth_prime_upper_bound(n).expect("Bound on nth prime exceeds usize"))
        .nth(n)
        .expect("nth prime was not in bounds")
}