prime-formula 0.3.1

//! Adds support for prime constellations
//
// Copyright: Adam Cottrell (cottrela@gmail.com)

use bitvec::prelude::*;
use rayon::prelude::*;
use rayon::slice::ParallelSliceMut;

use crate::common::{ConfigMode, PRIME_ROOTS, align_to_cycle};
use crate::sieve::{SieveMode, get_raw_composites};

/// Trait for filtering prime roots into constellation patterns
///
/// Implemented for the PRIME_ROOTS array to find indexes that form:
/// - Twin primes (difference 2)
/// - Prime triplets (differences 2/4 or 4/2)
/// - Higher-order constellations up to 6<sup>1</sup>
///
/// <sup>(1)</sup> *According to this novel theory, no constellations higher
/// than 6 can occur due to the arrangement of the prime roots.*
pub trait PrimeRootFilters {
    fn filter_twins(&self) -> Vec<u8>;
    fn filter_triplets(&self) -> Vec<u8>;
    fn filter_quadruplets(&self) -> Vec<u8>;
    fn filter_quintuplets(&self) -> Vec<u8>;
    fn filter_sextuplets(&self) -> Vec<u8>;
}

impl PrimeRootFilters for [u8; 48] {
    /// Find indexes of prime roots forming twin primes (p, p+2)
    fn filter_twins(&self) -> Vec<u8> {
        // Diff of 2
        self.windows(2)
            .enumerate()
            .filter_map(|(i, pair)| (pair[1] == pair[0] + 2).then_some(i as u8))
            .collect()
    }

    /// Find indexes of prime roots forming triplets (p, p+2, p+6) or (p, p+4, p+6)
    fn filter_triplets(&self) -> Vec<u8> {
        // Diff of 2 4 or 4 2
        self.windows(3)
            .enumerate()
            .filter_map(|(i, triple)| {
                let (a, b, c) = (triple[0], triple[1], triple[2]);
                let valid = a + 6 == c && (a + 2 == b || a + 4 == b);
                valid.then_some(i as u8) // get the first index
            })
            .collect()
    }

    /// Find indexes of prime roots forming quadruplets (p, p+2, p+6, p+8)
    fn filter_quadruplets(&self) -> Vec<u8> {
        // Diff of 2 4 2
        self.windows(4)
            .enumerate()
            .filter_map(|(i, quad)| {
                let (a, b, c, d) = (quad[0], quad[1], quad[2], quad[3]);
                let valid = a + 8 == d && a + 6 == c && a + 2 == b;
                valid.then_some(i as u8) // get the first index
            })
            .collect()
    }

    /// Find indexes of prime roots forming quintuplets, in the form:
    ///  (p, p+2, p+6, p+8, p+12) or (p, p+4, p+6, p+10, p+12)
    fn filter_quintuplets(&self) -> Vec<u8> {
        // Diff of 4 2 4 2 or 2 4 2 4
        self.windows(5)
            .enumerate()
            .filter_map(|(i, quint)| {
                let (a, b, c, d, e) = (quint[0], quint[1], quint[2], quint[3], quint[4]);
                let valid = a + 12 == e
                    && a + 6 == c
                    && ((a + 10 == d && a + 4 == b) || (a + 8 == d && a + 2 == b));
                valid.then_some(i as u8) // get the first index
            })
            .collect()
    }

    /// Find indexes of prime roots forming sextuplets,
    /// in the form (p, p+4, p+6, p+10, p+12, p+16)
    fn filter_sextuplets(&self) -> Vec<u8> {
        // Diff of 4 2 4 2 4
        self.windows(6)
            .enumerate()
            .filter_map(|(i, hex)| {
                let (a, b, c, d, e, f) = (hex[0], hex[1], hex[2], hex[3], hex[4], hex[5]);
                let valid = a + 16 == f && a + 12 == e && a + 10 == d && a + 6 == c && a + 4 == b;
                valid.then_some(i as u8) // get the first index
            })
            .collect()
    }
}

// Convenience function to expand constellation indices
fn expand_indices(indices: &[u8], count: u8) -> Vec<u8> {
    // Given a set of indices, expand them by count
    // and remove any duplicates
    // e.g. indices = [5, 9], count = 3
    // would return [5, 6, 7, 9, 10, 11]
    let mut unique_roots: Vec<u8> = indices
        .iter()
        .flat_map(|&i| (0..count).map(move |offset| i + offset))
        .collect();
    unique_roots.sort_unstable();
    unique_roots.dedup();
    unique_roots
}

/// Common helper function for finding prime constellations
fn find_constellation<T, F, G>(
    start: u128,
    end: u128,
    mode: &ConfigMode,
    filter_method: fn(&[u8; 48]) -> Vec<u8>,
    group_size: usize,
    initial_primes: &[T],
    create_tuple: F,
    include_initial: G,
) -> Vec<T>
where
    T: Ord + Send + Clone,
    F: Fn(u128, &[u128]) -> T + Sync + Send + Clone,
    G: Fn(&T) -> bool + Sync,
{
    let indices = filter_method(&PRIME_ROOTS);
    let mut unique_roots = expand_indices(&indices, group_size as u8);
    unique_roots.sort_unstable();
    unique_roots.dedup();

    let (b_start, b_end) = align_to_cycle(start, end);
    let composites = get_raw_composites(unique_roots, start, end, &SieveMode::Sieve);

    let mut primes: Vec<T> = indices
        .into_par_iter()
        .filter(|&i| (i as usize + group_size) <= PRIME_ROOTS.len())
        .flat_map(|i| {
            let roots: Vec<u128> = (0..group_size)
                .map(|offset| PRIME_ROOTS[i as usize + offset] as u128)
                .collect();

            // Need to clone function to ensure closure handled correctly
            let tuple_creator = create_tuple.clone();

            // Efficient bitwise combination
            let combined = (0..group_size)
                .map(|offset| composites[&(i + offset as u8)].as_bitslice())
                .fold(bitvec![0; (b_end - b_start) as usize], |acc, slice| {
                    acc | slice
                });

            combined
                .iter_zeros()
                .map(move |pos| {
                    let base = (b_start + pos as u128) * 210;
                    tuple_creator(base, &roots)
                })
                .collect::<Vec<_>>()
        })
        .collect();

    // Add initial primes with range check
    primes.extend(
        initial_primes
            .iter()
            .filter(|t| include_initial(t))
            .cloned(),
    );

    primes.par_sort_unstable();
    primes
}

/// Common helper function for counting prime constellations
fn count_constellation(
    start: u128,
    end: u128,
    mode: &ConfigMode,
    filter_method: fn(&[u8; 48]) -> Vec<u8>,
    group_size: usize,
) -> usize {
    let indices = filter_method(&PRIME_ROOTS);

    let mut unique_roots = expand_indices(&indices, group_size as u8);
    unique_roots.sort_unstable();
    unique_roots.dedup();

    let (b_start, b_end) = align_to_cycle(start, end);
    let composites = get_raw_composites(unique_roots, start, end, &SieveMode::Sieve);

    // Parallel count from sieve results
    let sieve_count: usize = indices
        .into_par_iter()
        .filter(|&i| (i as usize + group_size) <= PRIME_ROOTS.len())
        .map(|i| {
            (0..group_size)
                .map(|offset| composites[&(i + offset as u8)].as_bitslice())
                .fold(bitvec![0; (b_end - b_start) as usize], |acc, slice| {
                    acc | slice
                })
                .count_zeros()
        })
        .sum();

    // Add initial primes count with range check
    sieve_count
}

/// Get twin primes
pub fn get_twins(start: u128, end: u128, mode: &ConfigMode) -> Vec<(u128, u128)> {
    find_constellation(
        start,
        end,
        mode,
        PrimeRootFilters::filter_twins,
        2,
        &[(3, 5), (5, 7)],
        |base, roots| (base + roots[0], base + roots[1]),
        |&(p1, p2)| p1 >= start && p2 <= end,
    )
}

/// Get triplets
pub fn get_triplets(start: u128, end: u128, mode: &ConfigMode) -> Vec<(u128, u128, u128)> {
    find_constellation(
        start,
        end,
        mode,
        PrimeRootFilters::filter_triplets,
        3,
        &[(5, 7, 11), (7, 11, 13)],
        |base, roots| (base + roots[0], base + roots[1], base + roots[2]),
        |&(p1, _, p3)| p1 >= start && p3 <= end,
    )
}

/// Get quadruplets
pub fn get_quadruplets(start: u128, end: u128, mode: &ConfigMode) -> Vec<(u128, u128, u128, u128)> {
    find_constellation(
        start,
        end,
        mode,
        PrimeRootFilters::filter_quadruplets,
        4,
        &[(5, 7, 11, 13)],
        |base, roots| {
            (
                base + roots[0],
                base + roots[1],
                base + roots[2],
                base + roots[3],
            )
        },
        |&(p1, _, _, p4)| p1 >= start && p4 <= end,
    )
}

/// Get quintuplets
pub fn get_quintuplets(
    start: u128,
    end: u128,
    mode: &ConfigMode,
) -> Vec<(u128, u128, u128, u128, u128)> {
    find_constellation(
        start,
        end,
        mode,
        PrimeRootFilters::filter_quintuplets,
        5,
        &[(5, 7, 11, 13, 17), (7, 11, 13, 17, 19)],
        |base, roots| {
            (
                base + roots[0],
                base + roots[1],
                base + roots[2],
                base + roots[3],
                base + roots[4],
            )
        },
        |&(p1, _, _, _, p5)| p1 >= start && p5 <= end,
    )
}

/// Get sextuplets
pub fn get_sextuplets(
    start: u128,
    end: u128,
    mode: &ConfigMode,
) -> Vec<(u128, u128, u128, u128, u128, u128)> {
    find_constellation(
        start,
        end,
        mode,
        PrimeRootFilters::filter_sextuplets,
        6,
        &[(7, 11, 13, 17, 19, 23)],
        |base, roots| {
            (
                base + roots[0],
                base + roots[1],
                base + roots[2],
                base + roots[3],
                base + roots[4],
                base + roots[5],
            )
        },
        |&(p1, _, _, _, _, p6)| p1 >= start && p6 <= end,
    )
}

/// Count twin primes
pub fn count_twins(start: u128, end: u128, mode: &ConfigMode) -> usize {
    let mut count = count_constellation(start, end, mode, PrimeRootFilters::filter_twins, 2);

    count += [(3, 5), (5, 7)]
        .iter()
        .filter(|&&(p1, p2)| p1 >= start && p2 <= end)
        .count();

    count
}

/// Count triplets
pub fn count_triplets(start: u128, end: u128, mode: &ConfigMode) -> usize {
    let mut count = count_constellation(start, end, mode, PrimeRootFilters::filter_triplets, 3);

    count += [(5, 7, 11), (7, 11, 13)]
        .iter()
        .filter(|&&(p1, _, p3)| p1 >= start && p3 <= end)
        .count();
    count
}

/// Count quadruplet constellations
pub fn count_quadruplets(start: u128, end: u128, mode: &ConfigMode) -> usize {
    let mut count = count_constellation(start, end, mode, PrimeRootFilters::filter_quadruplets, 4);

    count += [(5, 7, 11, 13)]
        .iter()
        .filter(|&&(p1, _, _, p4)| p1 >= start && p4 <= end)
        .count();

    count
}

/// Count quintuplet constellations
pub fn count_quintuplets(start: u128, end: u128, mode: &ConfigMode) -> usize {
    let mut count = count_constellation(start, end, mode, PrimeRootFilters::filter_quintuplets, 5);

    count += [(5, 7, 11, 13, 17), (7, 11, 13, 17, 19)]
        .iter()
        .filter(|&&(p1, _, _, _, p5)| p1 >= start && p5 <= end)
        .count();
    count
}

/// Count sextuplet constellations
pub fn count_sextuplets(start: u128, end: u128, mode: &ConfigMode) -> usize {
    let mut count = count_constellation(start, end, mode, PrimeRootFilters::filter_sextuplets, 6);

    count += [(7, 11, 13, 17, 19, 23)]
        .iter()
        .filter(|&&(p1, _, _, _, _, p6)| p1 >= start && p6 <= end)
        .count();
    count
}