/* author Bill Wood */

use rayon::prelude::*;
use std::io::Write;

fn main() {
    let mut args = std::env::args()
        .skip(1)
        .map(|a| a
            .split('_')
            .fold("".to_string(), |a, s| a + s)
            .parse::<usize>());

    if args.len() <= 3 {
        if let (Ok(mut low), Ok(mut high), Ok(seg_size_kb), ) = (
                args.next().unwrap_or(Ok(1_000_000_000)),
                args.next().unwrap_or(Ok(2_000_000_000)),
                args.next().unwrap_or(Ok(64)),) {
            println!("S15: Finding sexy primes with args {} {} {}", low, high, seg_size_kb, );

            low = low | 1;   // lowest odd >= low
            if low < 3 {
                low = 3;
            }
            high = (high - 1) | 1;  // highest odd <= high

            // prepare for parallel segment sieve
            let primes_per_seg = seg_size_kb.next_power_of_two()*1024*8*2;
            let mut seg_low = low;
            let mut segs = vec![];
            while seg_low < high {
                let mut seg_high = seg_low + primes_per_seg - 2;
                if seg_high > high {
                    seg_high = high;
                }
                segs.push((seg_low, seg_high + 6));
                seg_low = seg_high + 2;
            }
            let lookup_table = PrimesIter::lookup_table_gen();

            // get sieving primes
            let sieve_primes = simple_sieve(high);
            // find sexy primes
            let (sexy_count, nprimes) =
                segs.into_par_iter()
                    .map(|(low, high)| {
                        let n = (high + 2 - low) >> 1;
                        let mut sieve_seg = Bvec::new(false, n);
                        for i in 0..sieve_primes.len() {
                            let p = sieve_primes[i] as usize;
                            let mut j = p*p;
                            if j > high {
                                break;
                            }
                            if j < low {
                                j += (low - 1 - j + p*2)/2/p*p*2;
                            }
                            j = (j - low) >> 1;
                            while j < n {
                                sieve_seg.set(j, true);
                                j += p;
                            }
                        }

                        let (mut sexy_count, mut nprimes) = (0_usize, 0_usize);
                        for p in PrimesIter::new(low, high, &sieve_seg, &lookup_table) {
                            if p + 6 <= high {
                                nprimes += 1;
                                if !sieve_seg.get(((p - low) >> 1) + 6/2) {
                                    sexy_count += 1;
                                }
                            }
                        }
                        (sexy_count, nprimes)
                    })
                   .reduce(|| (0, 0),
                           |a, b| (a.0 + b.0, a.1 + b.1));

            println!("{} of {} primes are sexy", sexy_count, nprimes);
            return;
        }
    }

    writeln!(
        std::io::stderr(),
        "Usage: sexy_primes [low [high [seg_size(kb)]]]"
    ).unwrap();
    std::process::exit(1);
}

fn simple_sieve(mut high: usize) -> Vec<u32> {
    // x where x: odd and (x+2)^2 > high and x^2 < high
    high = (sqrt(high) - 1) | 1;
    let mut pvec = vec![false; (high + 1)/2];
    let mut primes = vec![];
    let n = high >> 1;
    let mut i = 3;
    while i*i <= high {
        if !pvec[i >> 1] {
            primes.push(i as u32);
            let mut j = (i*i) >> 1;
            while j <= n {
                pvec[j] = true;
                j += i;
            }
        }
        i += 2;
    }
    i >>= 1;
    while i <= n {
        if !pvec[i] {
            primes.push(((i << 1) | 1) as u32);
        }
        i += 1;
    }
    primes
}

fn sqrt(n: usize) -> usize {
    let (mut r1, mut r) = (n, n + 1);
    while r1 < r {
        r = r1;
        r1 = (r1 + n/r1) >> 1
    }
    r
}

struct Bvec {
    v: Vec<u8>
}
impl Bvec {
    fn new(val: bool, s: usize) -> Self {
        let valu8 = if val { 255 } else { 0 };
        Bvec { v: vec![valu8; (s >> 3) + if (s & 7) == 0 { 0 } else { 1 } ] }
    }
    #[inline(always)]
    fn get(&self, index: usize) -> bool {
        self.v[index >> 3] & 1 << (index & 7) != 0
    }
    #[inline(always)]
    fn set(&mut self, mut index: usize, val: bool) {
        let bitpos = 1 << (index & 7);
        index >>= 3;
        if val {
            self.v[index] |= bitpos;
        } else {
            self.v[index] &= !bitpos;
        }
    }
}

struct PrimesIter<'a> {
    sieve_index: usize,
    base: usize,
    low: usize,
    high: usize,
    sieve: &'a [u8],
    lookup_table: &'a Vec<Vec<u8>>,
    lookup_table_cur: &'a Vec<u8>,
    lookup_table_cur_index: usize,
}
impl<'a> Iterator for PrimesIter<'a> {
    type Item = usize;
    #[inline(always)]
    fn next(&mut self) -> Option<Self::Item> {
        while self.lookup_table_cur_index >= self.lookup_table_cur.len() {
            if self.sieve_index >= self.sieve.len() {
                return None;
            }
            self.lookup_table_cur = &self.lookup_table[self.sieve[self.sieve_index] as usize];
            self.lookup_table_cur_index = 0;
            self.base = self.sieve_index*16 + self.low;
            self.sieve_index += 1;
        }
        let prime = self.base + self.lookup_table_cur[self.lookup_table_cur_index] as usize;
        self.lookup_table_cur_index += 1;
        if prime <= self.high {
            Some(prime)
        } else {
            None
        }
    }
}
impl<'a> PrimesIter<'a> {
    fn new(low: usize, high: usize, sieve: &'a Bvec, lookup_table: &'a Vec<Vec<u8>>) -> Self {
        PrimesIter { sieve_index: 0, base: 0, low, high, sieve: &sieve.v, lookup_table,
            lookup_table_cur: &lookup_table[0], lookup_table_cur_index: 8 }
    }
    fn lookup_table_gen() -> Vec<Vec<u8>> {
        let mut t = vec![vec![]; 256];
        for i in 0..256 {
            for j in 0..8 {
                if !i & 1 << j != 0 {
                    t[i].push((2*j) as u8);
                }
            }
        }
        t
    }
}