#!/usr/bin/env python3

import time
from bitarray import bitarray
from typing import List
from multiprocessing import Pool

PRIMES_FILE_NAME = "primes.txt"


class Eratosthenes:
    def __init__(self, process: int = 0):
        self.process = process

    @staticmethod
    def sequential_sieve(limit: int) -> List[int]:
        """
        Sequential version of eratosthenes sieve
        """

        result = bitarray(limit + 1)
        result.setall(1)
        result[0] = 0
        result[1] = 0
        for value in range(2, int(limit ** 0.5) + 1):
            if not result[value]:
                continue
            multiple = value + value
            while multiple <= limit:
                result[multiple] = 0
                multiple += value

        return [index for index, value in enumerate(result) if value]

    @staticmethod
    def worker(left_bound: int, right_bound: int) -> List[int]:
        """
        Worker executes sieving for a given range
        """

        left_bound = max(left_bound, 2)  # min value is 2
        print(f"worker range ({left_bound}, {right_bound})")

        primes = list(map(int, open(PRIMES_FILE_NAME).read().split()))
        result = list(range(left_bound, right_bound))
        prime_limit = int(right_bound ** 0.5)
        for prime in primes:
            if prime > prime_limit:
                break
            multiple = (-left_bound) % prime
            if multiple + left_bound == prime:
                multiple += prime
            size = right_bound - left_bound
            while multiple < size:
                result[multiple] = 0
                multiple += prime

        return [value for value in result if value]

    def parallel_sieve(self, limit: int) -> List[int]:
        """
        Parallel version of eratosthenes sieve
        """

        # find all primes up to sqrt(limit) and save it to file
        primes = self.sequential_sieve(int(limit ** 0.5))
        primes = ' '.join(map(str, primes))

        with open(PRIMES_FILE_NAME, "w") as file_handler:
            file_handler.write(primes)

        # parallel processing
        result = []
        with Pool(self.process) as pool:
            params = [(limit // self.process * x, limit // self.process * (x + 1))
                      for x in range(self.process)]
            for primes in pool.starmap(Eratosthenes.worker, params):
                result.extend(primes)

        return result

    def compute(self, limit: int) -> int:
        """
        This method returns the number of found primes up to 'limit'
        """

        # when number of processes is less than 2 do sequential processing
        if self.process < 2:
            print("Sequential processing...")
            return len(self.sequential_sieve(limit))

        print("Parallel processing...")
        return len(self.parallel_sieve(limit))


if __name__ == "__main__":
    sieve = Eratosthenes(10)
    start = time.time()
    print(f"num of primes: {sieve.compute(10**8)}")
    stop = time.time()
    print(f"time: {stop-start:.02f}")