//! Functions that sieve primes from the integers up to an input maximum.

use std::thread;
use std::sync::mpsc;
use std::error::Error;

extern crate itertools;

use itertools::Itertools;

/// Generates the sequence j = i^2, i^2+i, i^2+2i, i^2+3i, ...
///
/// This `Iterator` does not deal with overflow past usize; the caller must plan
/// accordingly.
struct SquareMultiple {
  curr: usize,
  increment: usize,
}

impl Iterator for SquareMultiple {
  type Item = usize;

  fn next(&mut self) -> Option<usize> {
    let val = self.curr;
    self.curr += self.increment;
    Some(val)
  }
}

impl SquareMultiple {
  fn new(term: usize) -> Self {
    SquareMultiple { curr: term * term, increment: term }
  }
}

/// An implementation of the sieve of Eratosthenes, as described in [the
/// Wikipedia article][wiki].
///
/// # Examples
///
/// Compute the primes in the first ten integers.
///
/// ```
/// assert_eq!(vec![2,3,5,7], prime_suspects::eratosthenes_sieve(10));
///
/// ```
///
/// Compute the last three primes up to 2^64.
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
///            prime_suspects::eratosthenes_sieve(65535)
///              .into_iter().rev().take(3)
///              .collect::<Vec<usize>>());
/// ```
///
/// [wiki]: https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Algorithm_and_variants
pub fn eratosthenes_sieve(max_val: usize) -> Vec<usize> {
  // Algorithm notes: The sieve works like so (Wikipedia pseudocode inlined):
  // Input: an integer n > 1
  // Let A be an array of Boolean values, indexed by integers 2 to n,
  // initially all set to true.
  let mut bool_vec = vec![true; max_val];
  // for i = 2, 3, 4, ..., not exceeding √n:
  let mut top_sieve = max_val as f64;
  // We have to add 1 because the sqrt coerced to an int is √floor(n)
  top_sieve = top_sieve.sqrt().ceil();
  
  for sieve_term in 2..(top_sieve as usize) {
    // if A[i] is true:
    if bool_vec[sieve_term] == true {
      // for j = i^2, i^2+i, i^2+2i, i^2+3i, ...,
      for j in SquareMultiple::new(sieve_term)
        .take_while(|&term| term < max_val) { // ...not exceeding n
          bool_vec[j] = false; // A[j] := false
        }
    }
  }
  // Output: all i such that A[i] is true.
  bool_vec[2..]
    .iter().enumerate()       // enumerate provides index
    .filter(|&(_, val)| *val) // filter out false
    .map(|(ind, _)| ind + 2)  // skips even numbers which we know are composite
    .collect()    
}

/// A [segmented approach][wiki] to sieveing, keeping memory use to O(√n). As
/// Sorensen states, this is the most practical optimization to the sieve of
/// Eratosthenes.
///
/// # Examples
///
/// Compute the last three primes up to 2^64 + 1. (The segmented sieve functions
/// call `eratosthenes_sieve` for lower values for reasons of efficiency.)
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
///            prime_suspects::segmented_sieve(65537, 256)
///              .into_iter().rev().take(3)
///              .collect::<Vec<usize>>());
/// ```
///
/// [wiki]: https://en.wikipedia.org/wiki/Sieve_of_Eratosthenes#Segmented_sieve
pub fn segmented_sieve(max_val: usize, segment_size: usize) -> Vec<usize> {
  if max_val <= 2usize.pow(16) {
    // early return if the highest value is small enough (empirical)
    return eratosthenes_sieve(max_val);
  }

  // if the segment size parameter is larger than alpha, the algorithm becomes
  // superlinear (see wiki), so we replace.
  let alpha = (max_val as f64).sqrt() as usize;
  let segment_size = if segment_size > alpha {
    println!("Segment size is larger than √{}. Reducing to {} to keep resource use down.",
             max_val, segment_size);
    alpha
  } else {
    segment_size
  };

  // get the primes up to the first alpha (even if the segment size is smaller)
  let small_primes = eratosthenes_sieve(alpha);
  let mut big_primes = small_primes.clone();

  // As Sorensen says, we need to construct a sequence over each segment, in
  // the interval [start + 1, start + segment_size] that begins with
  // (start + this_prime - ( start mod p)), and increases by p up to
  // (start + segment_size).
  // That sequence will be the values to sieve out of this_segment.
  for this_segment in (segment_size..max_val).collect::<Vec<_>>()
    .chunks(segment_size) {
    big_primes.extend(sieve_segment(&small_primes, this_segment));
  }
   
  big_primes
}

/// Same algorithm as the regular `segmented_sieve` function, but each larger
/// segment to be sieved is computed in parallel threads.
///
/// # Examples
///
/// Compute the last three primes up to 2^64 + 1. (The segmented sieve functions
/// call `eratosthenes_sieve` for lower values for reasons of efficiency.)
///
/// ```
/// assert_eq!(vec![65521, 65519, 65497],
///            prime_suspects::segmented_sieve_parallel(65537, 256)
///              .into_iter().rev().take(3)
///              .collect::<Vec<usize>>());
/// ```
pub fn segmented_sieve_parallel(max_val: usize, mut segment_size: usize) -> Vec<usize> {
  if max_val <= ((2 as i64).pow(16) as usize) {
    // early return if the highest value is small enough (empirical)
    return eratosthenes_sieve(max_val);
  }

  if segment_size > ((max_val as f64).sqrt() as usize) {
    segment_size = (max_val as f64).sqrt() as usize;
    println!("Segment size is larger than √{}. Reducing to {} to keep resource use down.",
             max_val, segment_size);
  }
  
  let small_primes = eratosthenes_sieve((max_val as f64).sqrt() as usize);
  let mut big_primes = small_primes.clone();

  let (tx, rx): (mpsc::Sender<Vec<usize>>, mpsc::Receiver<Vec<usize>>) = mpsc::channel();
  
  let extended_vec: Vec<_> = (segment_size..max_val).collect();
  let extended_segments = extended_vec.chunks(segment_size);
  for this_segment in extended_segments.clone() {
    let this_segment = this_segment.to_vec();
    let small_primes = small_primes.clone();
    let tx = tx.clone();

    thread::spawn(move || {
      let sieved_segment = sieve_segment(&small_primes, &this_segment);
      tx.send(sieved_segment).unwrap_or_else(|why| {
        println!("Segment sieve failed: {}", why.description());
      });
    });
  }

  for _ in 0..extended_segments.count() {
    big_primes.extend(&rx.recv().unwrap());
  }

  big_primes
}

// The actual function that does the sieving on a segment. It takes a vector of
// primes and vector of primes expected to be larger than any of the primes.
// (The function doesn't test for that condition.)
fn sieve_segment(primes: &Vec<usize>, target_vec: &[usize]) -> Vec<usize> {
  let mut sieved_segment = Vec::with_capacity(target_vec.len());
  sieved_segment.extend(target_vec);
  for &this_prime in primes {
    if !sieved_segment.is_empty() {
      let first_val = target_vec[0];
      let mut starting_offset = first_val % this_prime;
      starting_offset = if starting_offset == 0 { this_prime } else { starting_offset };

      let first_val = first_val + this_prime - starting_offset;
      let last_val: &usize = target_vec.last().unwrap();

      // hack for inclusive range while RFC is figured out. see
      // https://www.reddit.com/r/rust/comments/3xkfro/what_happened_to_inclusive_ranges/
      let sieve_vec: Vec<_> = (first_val..(*last_val + 1))
        .step(this_prime)
        .collect();
      
      sieved_segment = sieved_segment
        .iter()
        .filter(|check_num| !sieve_vec.contains(check_num))
        .cloned()
        .collect();
    }
  }

  sieved_segment
}