# PrimeCounter.py

# Python has multi-threading, but that doesn't take advantage of multiple cores because
# of an underlying global interpreter lock (GIL)

# multiprocessing is heavier (a new process is created). 

# the code below generates the correct answer for the threaded code, but only one processor is used
# It does run multiple processes if we use multithreading, but getCount returns 0 at the end. I think
# we would need to use a multiprocessing.Array to get information back to the main process.

# In order to make it possible to return a value to the main code with multiprocessing, I needed to use 
# the Value class, so this code is not really set up to work for both threading and multiprocessing.

from multiprocessing import Process as Worker
from multiprocessing import Value
import time

class PrimeCounter(Worker):
    
    # typedef struct {
    #     int id;
    #     int N;
    #     int chunk_size;
    #     int num_threads;
    #     int count;
    # } ThreadInfo;


    def __init__ (self, id, N, chunk_size, num_threads, shared_count):
        super().__init__() 
        self.id = id
        self.N = N
        self.chunk_size = chunk_size
        self.num_threads = num_threads
        self.count = 0
        self.shared_count = shared_count # this one is visible after the process is over. 

    def is_prime (self, n):
        if (n < 2):
            return False
        if (n == 2):
            return True
        if (n % 2 == 0):
            return False
        for i in range (3, n, 2):
            if (n % i == 0):
                return False;
        return True                    
        
    def run (self):
        print( "running thread with pid", self.pid )

#       int start = ti->chunk_size * ti->id;
#         for (int i = start; i < ti->N; i = i + ti->chunk_size*ti->num_threads ) {
#             for (int j = 0; j < ti->chunk_size; j++) {
        start = self.chunk_size*self.id
        for i in range( start, self.N, self.chunk_size*self.num_threads):
            for j in range(0, self.chunk_size ):
                if self.is_prime(i+j):
                    self.count += 1      
#         print( "found",self.count,"primes", self.getCount() )
        self.shared_count.value = self.count
        return 0       
    
    def getCount (self):
        return self.shared_count.value
        

def main ():
    N = 100000
    NUM_THREADS = 20

    pc = []
    return_values = []
    for i in range(NUM_THREADS):
        return_values.append( Value('i') )
        pc.append( PrimeCounter( id=i, N=N, chunk_size=100, num_threads=NUM_THREADS, 
                 shared_count=return_values[i]) )

    return_value0 = Value('i')
    pc0 = PrimeCounter( id=0, N=N, chunk_size=N, num_threads=1, shared_count=return_value0)

    t1 = time.time()
    
    for i in range(NUM_THREADS):
        pc[i].start()
       
    for i in range(NUM_THREADS):
        pc[i].join()
    
    count = 0
    for i in range(NUM_THREADS):
        count += pc[i].getCount()
        
    t2 = time.time()
    
    print("multi-thread count ", count, " primes")
    print("time is ", t2 - t1)
    
    t1 = time.time()
    
    pc0.start()
    pc0.join()
    
    t2 = time.time()
    
    print("1 thread count ", pc0.getCount(), " primes")
    print("time is ", t2 - t1)
    
    
if __name__ == "__main__":
    main()