import numpy as np
from mpi4py import MPI
from time import time
import matplotlib.pyplot as plt

def compute_pi(ns): 
    """
    input:  ns   list of block sizes
    output: pis  len(ns) estimated pi values
    method: pick n x n random points in a square area [-1,1] surrounding a circle radius 1
            probability of hitting the circle is
            P = (area of circle)/(area of square) = pi/(2*2) = hits/n => pi = 4*hits/n
    """
    pis = np.zeros(len(ns))
    for i,n in enumerate(ns):
        points = 2*np.random.random((n,2)) - 1 # random points in square
        # count points inside the unit circle; 1 (True) for hits, 0 (False) for misses
        hits = np.sum(np.linalg.norm(points,axis=-1) < 1)
        pis[i] = 4*hits/n        
    return pis

def taskdivider(N,nb,nproc):
    """
    input:  N     number of tasks
            nb    number of blocks
            nproc number of processes used for computation
    output: ns    (nproc,nb) array, N tasks divided as evenly as possible to nproc processes and nb blocks
                  low rank processes may get a few more tasks than the rest 
    """    
    # // gives integer result, discard remainder    
    ns = np.array([N//nproc//nb]*nb*nproc)
    # distribute remaining jobs evenly among blocks
    i = 0
    while ns.sum()<N:
        ns[i] += 1
        i += 1
    ns = ns.reshape(nproc,nb)
    return ns

def make_plots():
    timings, pis, stds = zip(*res)
    # pi estimates
    plt.figure(1)
    plt.errorbar(Ns,pis,stds,marker='o',ls='none',markersize=4,capsize=3)
    plt.xlabel('N')
    plt.xscale('log', base=2) # N in log2 scale
    plt.ylabel(r'estimate of $\pi$')
    #fname = f'pi_{nproc}.png'
    #print('saving pi estimate figure to file',fname)
    #plt.savefig(fname)
    
    # timings 
    plt.figure(2)
    plt.plot(Ns,timings,'o-')
    plt.xlabel('N')
    plt.xscale('log', base=2) # N in log2 scale
    plt.ylabel('timings (s)')
    #fname = f'timings_{nproc}.png'
    #print('saving timings figure to file',fname)
    #plt.savefig(fname)
    plt.show()

if __name__=='__main__':

    # init MPI
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    nproc = comm.Get_size()     
    if rank==0: print(nproc,'processes')

    # parameter: upper limit for N
    N_max = 1e10
    nb = 10 #  target number of blocks (error estimation over blocks)

    # numbers of random points to be used in estimation of pi
    i_max = int(np.log2(N_max))+1
    Ns = [2**i for i in range(10,i_max,2)]
    
    if __debug__ and rank==0: print('Ns = ',Ns)        
    
    res = []
    for N in Ns:

        # find how many points each process need to compute
        ns = None
        if rank==0: ns = taskdivider(N,nb,nproc)            

        # distribute the numbers of points to processes
        ns = comm.scatter(ns, root=0)

        # run with -O to suppress this printout
        if __debug__:
            print(f'rank {rank} will compute blocks with number of points {ns}, in total {np.sum(ns)} points')

        if rank==0: tic = time()
        # 
        # calculate pi estimates
        pi_est = compute_pi(ns)
        #
        # gather results from processes
        pis = comm.gather(pi_est, root=0)

        if rank==0:
            toc = time()
            pi  = np.mean(pis)
            std = np.std(pis)
            print(f'N = {N:<10d}  pi = {pi:<30.20f} +/- {std:<30.20f}  timing {toc-tic:<.4f} s')
            res.append([toc-tic,pi,std])

    # plot results
    if rank==0: make_plots()