from rmtkernel import ensemble_diffs, normalize_diffs, GOE
import numpy as np
from IPython.parallel import Client

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

def generate_wigner_data():
    s = np.linspace(0.0,4.0,400)
    rhos = wigner_dist(s)
    return s, rhos

s, rhos = generate_wigner_data()

plot(s, rhos)
xlabel('Normalized level spacing s')
ylabel('Probability $\rho(s)$')

def serial_diffs(num, N):
    """Compute the nearest neighbor distribution for num NxX matrices."""
    diffs = ensemble_diffs(num, N)
    normalized_diffs = normalize_diffs(diffs)
    return normalized_diffs

serial_nmats = 1000
serial_matsize = 50

%timeit -r1 -n1 serial_diffs(serial_nmats, serial_matsize)

serial_diffs = serial_diffs(serial_nmats, serial_matsize)

hist_data = hist(serial_diffs, bins=30, normed=True)
plot(s, rhos)
xlabel('Normalized level spacing s')
ylabel('Probability $P(s)$')

def parallel_diffs(rc, num, N):
    nengines = len(rc.targets)
    num_per_engine = num/nengines
    print "Running with", num_per_engine, "per engine."
    ar = rc.apply_async(ensemble_diffs, num_per_engine, N)
    diffs = np.array(ar.get()).flatten()
    normalized_diffs = normalize_diffs(diffs)
    return normalized_diffs

client = Client()
view = client[:]
view.run('rmtkernel.py')
view.block = False

parallel_nmats = 40*serial_nmats
parallel_matsize = 50

%timeit -r1 -n1 parallel_diffs(view, parallel_nmats, parallel_matsize)