#-------------------------------------------------------------------------------
# Imports
#-------------------------------------------------------------------------------

import time
import numpy

# import IPython.kernel.magic
from IPython.parallel import Client, Reference
from IPython.parallel import error

rc = Client()
mux = rc[:]

#-------------------------------------------------------------------------------
# Setup
#-------------------------------------------------------------------------------

%load_ext parallelmagic

mux.block = True
mux.clear()
mux.activate()

n = len(rc)
assert n >= 4, "Not Enough Engines: %i, 4 needed for this script"%n

values = [
    10,
    1.0,
    range(100),
    ('asdf', 1000),
    {'a': 10, 'b': 20}
    ]

keys = ['a','b','c','d','e']

sequences = [
    range(100),
    numpy.arange(100)
]

#-------------------------------------------------------------------------------
# Blocking execution
#-------------------------------------------------------------------------------

# Execute

with mux.sync_imports():
    import math
    import numpy

mux.execute('a = 2.0*math.pi')
print mux['a']

for id in mux.targets:
    mux.execute('b=%d' % id, targets=id)


try:
    mux.execute('b = 10',targets=max(mux.targets)+1)
except IndexError:
    print "Caught invalid engine ID OK."

try:
    mux.apply(lambda : 1/0)
except error.CompositeError:
    print "Caught 1/0 correctly."



%px print a, b

try:
    %px 1/0
except error.CompositeError:
    print "Caught 1/0 correctly."


# %autopx

%px a = numpy.random.random((4,4))
%px a = a+a.transpose()
# %autopx

print mux.apply(numpy.linalg.eigvals, Reference('a'))



mux.targets = [0,2]
%px a = 5
mux.targets = [1,3]
%px a = 10
mux.targets = rc.ids
print mux['a']


# Push/Pull

mux.push(dict(a=10, b=30, c={'f':range(10)}))
mux.pull(('a', 'b'))

for id in mux.targets:
    mux.push(dict(a=id), targets=id)


for id in mux.targets:
    mux.pull('a', targets=id)

mux.pull('a')


mux['a'] = 100
mux['a']

# get_result/reset/keys

mux.get_result()
%result
mux.apply(lambda : globals().keys())
mux.clear()
mux.apply(lambda : globals().keys())

try:
    %result
except error.CompositeError:
    print "Caught IndexError ok."


%px a = 5
mux.get_result(-1)
mux.apply(lambda : globals().keys())

# Queue management methods

%px import time
ars = [mux.apply_async(time.sleep, 2.0) for x in range(5)]


mux.queue_status()
time.sleep(3.0)
mux.abort(ars, block=True)
mux.queue_status()
time.sleep(2.0)
mux.queue_status()

mux.wait(ars)

for ar in ars:
    try:
        ar.r
    except error.CompositeError:
        print "Caught QueueCleared OK."


# scatter/gather

mux.scatter('a', range(10))
mux.gather('a')
mux.scatter('b', numpy.arange(10))
mux.gather('b')

#-------------------------------------------------------------------------------
# Non-Blocking execution
#-------------------------------------------------------------------------------

mux.block = False

# execute

ar1 = mux.execute('a=5')
with mux.sync_imports():
    import sets

ar1.wait()

ar1 = mux.execute('1/0')
ar2 = mux.execute('c = sets.Set()')

mux.wait((ar1, ar2))
try:
    ar1.r
except error.CompositeError:
    print "Caught ZeroDivisionError OK."

ar = mux.execute("arint 'hi'")
ar.r

ar = mux.apply(lambda : 1/0)
try:
    ar.r
except error.CompositeError:
    print "Caught ZeroDivisionError OK."

# Make sure we can reraise it!
try:
    ar.r
except error.CompositeError:
    print "Caught ZeroDivisionError OK."

# push/pull

ar1 = mux.push(dict(a=10))
ar1.get()
ar2 = mux.pull('a')
ar2.r


# This is a command to make sure the end of the file is happy.

print "The tests are done!"