'''
==========================================================================
Python for Parallelism in Introductory Computer Science Education
SC '13 HPC Educators Program
Steven Bogaerts, Wittenberg University
Joshua Stough, Washington and Lee
http://www.joshuastough.com/SC13
MIT License: see README_LICENSE.txt


file: parallelSumPrimes.py
author: stough
Summary: Sums the primes up through n, both sequentially and
in parallel.

Uses Pool/map paradigm, where we evenly divide the independent
work and have each process of the pool compute the partial
result.  Then sum up all the results.

Here, we also split the work many more times than the
number of processes.  that way, the work comes in smaller
chunks and is therefore more likely to be distributed evenly.

$ python[3] parallelSumPrimes.py [1000000]

See:
http://docs.python.org/library/multiprocessing.html
http://docs.python.org/py3k/library/multiprocessing.html
==========================================================================
'''

import math, time, sys
from multiprocessing import Pool, cpu_count

#Dependencies defined below main()

def main():
    """
    The main function, which times both sequential and
    parallel versions of sumPrimes, computing the sum of
    all primes up through the argument n.

    """
    n = 1000000
    if len(sys.argv) > 1:
        n = int(sys.argv[1])

    #Sequential version.
    start = time.time()
    resSeq = sumPrimes(2,n)
    elapsed = time.time() - start
    print('Sequential sumPrimes took %f sec.' % (elapsed))

    #So that cpu usage shows a lull.
    time.sleep(3)

    #Parallel version.
    start = time.time()
    cpus = cpu_count() 
    
    args = buildSegs([int(x) for x in linspace(2,n, 128)])
    """While the cpu count would be an obvious choice for 
       splitting up the problem, isPrime takes longer on 
       larger numbers, so we'll split it up many times, so
       that we more evenly distribute the load."""

    #Create a pool of processors.
    pool = Pool(processes=cpus)
    #Each process executes wrapSumPrimes on a particular args[i],
    #tuple expressing a low and high integer to consider.
    results = pool.map(wrapSumPrimes, args)
    resPar = sum(results)

    elapsed = time.time() - start
    print('Parallel sumPrimes took %f sec.' % (elapsed))

    if resSeq != resPar:
        print('Uh-oh: the two sums are not equal: %d, %d' \
            % (resSeq, resPar))

        

def isPrime(n):
    """Returns True if n is prime, False otherwise."""
    if n == 2:
        return True
    elif n % 2 == 0:
        return False

    for i in range(3, int(math.ceil(math.sqrt(n))) + 1, 2):
        if n % i == 0:
            return False

    return True
    

def sumPrimes(a,b):
    """Computes the sum of the primes from a to b"""
    return sum([x for x in range(a,b+1) if isPrime(x)])

def wrapSumPrimes(ab):
    """
    The wrap is necessary in Python 3 because the 
    function sent to the processes in a pool must 
    accept only one argument, and argument unpacking
    in the method header is deprecated in Python 3.
    """
    (a,b) = ab
    return sumPrimes(a,b)

def linspace(a,b,nsteps):
    """returns list of simple linear steps from a to b in nsteps."""
    ssize = float(b-a)/(nsteps-1)
    return [a + i*ssize for i in range(nsteps)]

def buildSegs(lyst):
    """Builds the problem segments for sumPrimes."""
    args = []
    for i in range(len(lyst)-1):
        args.append((lyst[i], lyst[i+1]))
    
    for i in range(1, len(args)):
        (a,b) = args[i]
        args[i] = (a+1,b)
    return args


#Call the main method if run from the command line.
if __name__ == '__main__':
    main()


"""
Alternative to the kind-of complicated buildSegs, but it's
a little slower:
args = [2] + list(range(3, n+1, 2))
pool = Pool(processes=cpus)
results = pool.map(isPrime, args)
#because the results come out in order...
resPar = sum([a*b for a,b in zip(args, [(0,1)[x] for x in results])])
"""