####################################################
#  Calculate Mandelbrot set and save it as a bmp image
#
#  Parallel version uses master-worker approach,
#  with master node at rank zero; the work units
#  distributed to workers consist of a single
#  image row. It needs at least two MPI nodes to work.
#
####################################################

import mpi
import bmp

# maximal number of iterations to compute a pixel
MAX_ITER = 256

# pixel computation function
def pixel(c):
    z = 0
    for i in range(MAX_ITER):
        z = z*z+c
        if abs(z) >= 2.0:
            return i
    return MAX_ITER

# image dimensions
nx = 1024
ny = 1024

workers = []
for i in range(mpi.size):
    workers.append(0)

if mpi.rank == 0:
    # "master" node:

    # number of jobs being processed and the next image row to compute
    jobs = 0
    row  = 0

    # initialize list of image rows
    image = []
    for i in range(ny):
        image.append(-1)

    # assign initial work units to all workers
    for n in range(1, mpi.size):
        mpi.send(row, n)
        row += 1
        jobs += 1
        
    # master's main loop:
    while jobs > 0:
        # receive computed result from any worker
        result, status = mpi.recv()

        workers[status.source] += 1

        # incorporate newly computed row into image data
        image[result[0]] = result[1]

        if row < ny:
            # send new work unit to the (now) idle worker
            mpi.send(row, status.source)
            row += 1
        else:
            # no jobs remain: dismiss the worker
            mpi.send(-1, status.source)
            jobs -= 1

    # convert data to color image and save it in a file
    bmp.write_image('image.bmp', nx, ny, image, MAX_ITER)
    for w in range(1,mpi.size):
        print workers[w],"tasks completed by worker",w
    
else:
    # "worker" node:
    while 1:
        # receive work unit info
        row, status = mpi.recv()
        # check if we're still needed
        if row == -1:
            break
        # compute row of image
        rdata = []

        # Magic here... 4.0j is a complex number
        c = 4.0j*row/ny-2.0j
        for x in range(nx):
            rdata += [pixel(c+4.0*x/nx-2.0)]
        # send the result to master
        mpi.send([row, rdata], 0)