#-------------------------------------------------------------------------------
# Driver code that the client runs.
#-------------------------------------------------------------------------------
# To run this code start a controller and engines using:
# ipcluster -n 2
# Then run the scripts by doing irunner rmt.ipy or by starting ipython and
# doing run rmt.ipy.

from rmtkernel import *
from IPython.kernel import client


def wignerDistribution(s):
    """Returns (s, rho(s)) for the Wigner GOE distribution."""
    return (numpy.pi*s/2.0) * numpy.exp(-numpy.pi*s**2/4.)


def generateWignerData():
    s = numpy.linspace(0.0,4.0,400)
    rhos = wignerDistribution(s)
    return s, rhos
    

def serialDiffs(num, N):
    diffs = ensembleDiffs(num, N)
    normalizedDiffs = normalizeDiffs(diffs)
    return normalizedDiffs


def parallelDiffs(rc, num, N):
    nengines = len(rc.get_ids())
    num_per_engine = num/nengines
    print "Running with", num_per_engine, "per engine."
    rc.push(dict(num_per_engine=num_per_engine, N=N))
    rc.execute('diffs = ensembleDiffs(num_per_engine, N)')
	# gather blocks always for now
    pr = rc.gather('diffs')
    return pr.r


# Main code
if __name__ == '__main__':
    rc = client.MultiEngineClient()
    print "Distributing code to engines..."
    r = rc.run('rmtkernel.py')
    rc.block = False

    # Simulation parameters
    nmats = 100
    matsize = 30
    
    %timeit -n1 -r1 serialDiffs(nmats,matsize)
    %timeit -n1 -r1 parallelDiffs(rc, nmats, matsize)

    # Uncomment these to plot the histogram
    # import pylab
    # pylab.hist(parallelDiffs(rc,matsize,matsize))