import math
import time
from typing import Tuple

import numpy as np

from src import get_cache_size
from src.multiples import Multiples, Multiple, factory
from src.multiples import factory as multiples_factory
from src.primes import Primes
from src.primes import factory as primes_factory
from src.primesgen import Sieve


class SegmentedEratosthenes(Sieve):
    """
    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 _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 primes(self, in_max_val) -> list[int]:
        self._multiples = factory(in_max_val)
        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(math.sqrt(self._max_val)) + 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)
        return self._primes.all()

    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) 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.add(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.add(Multiple(p_2_val - seg_offset, p_val, max_val_seg))


if __name__ == '__main__':
    max_val = 1_000_000_000
    # print(f'\torg found all {len(Eratosthenes().primes(max_val)):,} primes <= {max_val:,}')
    # print(f'\tidx found all {len(IndexedEratosthenes().primes(max_val)):,} primes <= {max_val:,}')
    primes = primes_factory()  # if we supply a queue, primes will be broadcast instead of returning vals (and doesn't store them after they are sent)
    multiples = multiples_factory(max_val)
    segment_max_size = get_cache_size()
    start = time.time()
    prime_values = SegmentedEratosthenes(segment_max_size, primes, multiples).primes(max_val)
    duration = time.time() - start
    print(f'\tseg found all {len(prime_values):,} primes <= {max_val:,} in {duration:.2f} seconds')
