from typing import Iterable, Tuple

import numpy as np

from src.multiples import Multiples, Multiple
from src.primes import Primes


class Eratosthenes:
    """
    Classic Eratosthenes, leveraging numpy
    """

    def primes(self, in_max_val: int) -> list[int]:
        return list(self.primes_generator(in_max_val))

    def primes_generator(self, in_max_val: int) -> Iterable[int]:
        if in_max_val < 3:
            return []
        s = np.ones(in_max_val + 1, dtype=bool)
        s[1] = s[0::2] = False
        for p_2, p_t_2 in ((i ** 2, 2 * i) for i in (i for i in range(3, int(in_max_val ** 0.5) + 1, 2) if s[i])):
            s[p_2::p_t_2] = False
        return np.where(s)[0].astype(int)


class SegmentedEratosthenes(Eratosthenes):
    """
    Utilizes segmentation to handle numbers beyond the scope of the normal erat by reducing the mem used for the sieve

    The prime values can be broadcast, if a queue is provided, thus reducing the primes in mem
    please see the factory function in primes

    The multiples can be stored in mem, or offloaded to disk as needed
    please see the factory function in multiples

    Essentially a copy of Kim Walisch's C++ impl (which is way way way faster)
    https://github.com/kimwalisch/primesieve/wiki/Segmented-sieve-of-Eratosthenes

    but the off loading of the data to queues and/or disk is something I want to investigate more
    """

    def __init__(self, in_segment_max_size, in_primes: Primes, in_multiples: Multiples) -> None:
        self._seg_max_size = in_segment_max_size
        self._multiples = in_multiples
        self._primes = in_primes
        self._max_sqrt_val = None
        self._max_val = None

    def primes(self, in_max_val: int) -> list[int]:
        return list(self.primes_generator(in_max_val))

    def primes_generator(self, in_max_val) -> Iterable[int]:
        if in_max_val < 3:
            return []
        # prep data for all segments based on input
        self._max_val = in_max_val
        self._max_sqrt_val = int(self._max_val ** .5) + 1
        num_segs = int((self._max_val + 1) / self._seg_max_size) + 1
        if self._seg_max_size == self._max_val + 1:
            num_segs = 1

        # sieve each segment
        for seg_num in range(1, num_segs + 1):
            self._sieve_segment(seg_num)
        list(self._multiples)
        return self._primes.__iter__()

    def _create_sieve(self, seg_max_val: int, seg_offset: int, seg_num: int) -> np.ndarray[Tuple[int]]:
        new_sieve = np.ones(seg_max_val + 1, dtype=bool)
        # bye bye evens
        new_sieve[seg_offset % 2::2] = False
        if seg_num == 1:
            new_sieve[1] = False
        for m in self._multiples:
            new_sieve[m.start_val::m.p_val] = False
        self._multiples.advance(seg_max_val)
        return new_sieve

    def _sieve_segment(self, seg_num: int) -> None:
        seg_offset = (seg_num - 1) * self._seg_max_size
        max_val_seg = min(self._max_val - seg_offset, self._seg_max_size - 1)
        max_sqrt_val_seg = min(max(0, self._max_sqrt_val - seg_offset), self._seg_max_size)
        sieve = self._create_sieve(max_val_seg, seg_offset, seg_num)

        # erat (0 --> max_sqrt_val_seg)
        for p_val in (true_index + seg_offset for true_index in range(0, max_sqrt_val_seg - 1) if sieve[true_index]):
            self._primes.append(p_val)
            p_2_val = p_val * p_val
            p_2_val_seg = p_2_val - seg_offset
            self._multiples.append(Multiple(p_2_val_seg, p_val, max_val_seg))
            sieve[p_2_val_seg::p_val] = False

        # rest (max_sqrt_val_seg -- > end          and        additional offset of max_sqrt_val_seg)
        for p_val in (true_index + max_sqrt_val_seg + seg_offset for true_index in
                      np.where(sieve[max_sqrt_val_seg::])[0]):
            self._primes.append(p_val)
            p_2_val = p_val * p_val
            self._multiples.append(Multiple(p_2_val - seg_offset, p_val, max_val_seg))
