import pandas as pd

from IPython.parallel import Client

c = Client()

view = c[:]
view.block = True
view.activate() # enable magics

%px import fipy as fp
%px procID = fp.parallelComm.procID
%px Nproc = fp.parallelComm.Nproc

print view['procID']
print view['Nproc']

def setup(N):
    import fipy as fp
    import numpy as np
    np.random.seed(1)
    L = 1.
    m = fp.GmshGrid3D(nx=N, ny=N, nz=N, dx=L / N, dy=L / N, dz=L / N)
    v0 = 1.
    for dim in range(3):
        x = np.linspace(0., L, N)
        fx = np.sum(np.array([np.sin(2 * x * np.pi * i * np.random.random()) / i for i in range(1, N)]), axis=0)
        v0 = np.outer(v0, fx).flatten()
    v = fp.CellVariable(mesh=m)
    v0 = np.resize(v0, len(v)) ## Gmsh doesn't always give us the correct sized grid!
    eqn = fp.TransientTerm(1e-3) == fp.DiffusionTerm()
    v[:] = v0.copy()
    return eqn, v, v0

view['setup'] = setup

%%px

eqn, v, v0 = setup(60)
eqn.solve(v, dt=1.)
%timeit v[:] = v0.copy(); eqn.solve(v, dt=1.)

eqn, v, v0 = setup(60)
eqn.solve(v, dt=1.)
%timeit v[:] = v0.copy(); eqn.solve(v, dt=1.)

%px print eqn.getDefaultSolver().__class__
print eqn.getDefaultSolver().__class__

import fipy.solvers.pysparse as solvers
solver = solvers.linearPCGSolver.LinearPCGSolver(precon=None, iterations=500, tolerance=1e-15)

eqn.solve(v, solver=solver, dt=1.)
%timeit v[:] = v0.copy(); eqn.solve(v, solver=solver, dt=1.)

import fipy.solvers.trilinos as solvers
solver = solvers.linearPCGSolver.LinearPCGSolver(precon=None, iterations=500, tolerance=1e-15)

eqn.solve(v, solver=solver, dt=1.)
%timeit v[:] = v0.copy(); eqn.solve(v, solver=solver, dt=1.)

%%px
import fipy.solvers.trilinos as solvers
solver = solvers.linearPCGSolver.LinearPCGSolver(precon=None, iterations=500, tolerance=1e-15)
eqn.solve(v, solver=solver, dt=1.)
%timeit v[:] = v0.copy(); eqn.solve(v, dt=1., solver=solver)

setup_str = '''
import fipy as fp
import numpy as np
np.random.seed(1)
L = 1.
N = {N:d}
m = fp.GmshGrid3D(nx=N, ny=N, nz=N, dx=L / N, dy=L / N, dz=L / N)
v0 = 1.
for dim in range(3):
    x = np.linspace(0., L, N)
    fx = np.sum(np.array([np.sin(2 * x * np.pi * i * np.random.random()) / i for i in range(1, N)]), axis=0)
    v0 = np.outer(v0, fx).flatten()
v = fp.CellVariable(mesh=m)
v0 = np.resize(v0, len(v)) ## Gmsh doesn't always give us the correct sized grid!
eqn = fp.TransientTerm(1e-3) == fp.DiffusionTerm()
v[:] = v0.copy()

import fipy.solvers.{suite} as solvers
solver = solvers.linearPCGSolver.LinearPCGSolver(precon=None, iterations={iterations}, tolerance=1e-100)

eqn.solve(v, dt=1., solver=solver)
v[:] = v0.copy()
'''

timeit_str = '''
eqn.solve(v, dt=1., solver=solver)
fp.parallelComm.Barrier()
'''

import timeit
import itertools
%px import timeit

attempts = 3

view['timeit_str'] = timeit_str
view['attempts'] = attempts
view['setup_str'] = setup_str

runtimes = pd.DataFrame(columns=['iterations', 'N', 'mode', 'suite', 'run times'])

Ns = (10, 20, 30, 40, 50, 60, 70, 80)
iterations_ = (100, 200, 400, 800)
modes = ('serial', 'parallel')
suites = ('pysparse', 'trilinos')
                      
for N, iterations, mode, suite in itertools.product(Ns, iterations_, modes, suites):
    if mode == 'serial':
        timer = timeit.Timer(timeit_str, setup=setup_str.format(N=N, suite=suite, iterations=iterations))
        times = timer.repeat(attempts, 1)
        runtime = min(times)
    elif suite == 'trilinos':
        view['N'] = N
        view['iterations'] = iterations
        view['suite'] = suite
        %px timer = timeit.Timer(timeit_str, setup=setup_str.format(N=N, suite=suite, iterations=iterations))
        %px times = timer.repeat(attempts, 1)
        runtime = max(np.array(view['times']).min(axis=1))
    else:
        runtime = None
    
    runtimes = runtimes.append({'N' : N, 'iterations' : iterations, 'mode' : mode, 'suite' : suite, 'run time' : runtime})

runtimes.to_csv('laptop_runtimes.csv')

import pandas as pd
df = pd.read_csv('laptop_runtimes.csv', index_col=0)
iterations_ = (100, 200, 400, 800)

f = plt.figure(figsize=(15,5))
ax1 = f.add_subplot(131)
ax2 = f.add_subplot(132)
ax3 = f.add_subplot(133)

for iters in iterations_:
    parallel_df = df[(df['iterations'] == iters) & (df['mode'] == 'parallel') & (df['suite'] == 'trilinos')].sort('N')
    serial_df = df[(df['iterations'] == iters) & (df['mode'] == 'serial') & (df['suite'] == 'trilinos')].sort('N')
    pysparse_df = df[(df['iterations'] == iters) & (df['mode'] == 'serial') & (df['suite'] == 'pysparse')].sort('N')
    Ns = parallel_df['N']
    parallelTimes = np.array(parallel_df['run times'])
    serialTimes = np.array(serial_df['run times'])
    pysparseTimes = np.array(pysparse_df['run times'])
    ax1.semilogx(Ns**3, serialTimes / parallelTimes, label=iters)
    ax2.semilogx(Ns**3, pysparseTimes / serialTimes, label=iters)
    if iters == 800:
        ax3.loglog(Ns**3, parallelTimes, label='parallel Trilios')
        ax3.loglog(Ns**3, serialTimes, label='serial Trilinos')
        ax3.loglog(Ns**3, pysparseTimes, label='PySparse')
    
ax1.set_ylabel("Speed Up")
ax2.set_ylabel("Speed Up")
ax1.set_xlabel("$N^3$")
ax2.set_xlabel("$N^3$")
ax1.set_title("Speed up, serial / parallel")
ax2.set_title("Speed up, PySparse / Trilinos (serial)")
ax3.set_title("Raw Times, 800 iterations")
ax3.set_xlabel("$N^3$")
ax3.set_ylabel("Time (s)")

ax1.legend(loc='upper left')
ax2.legend(loc='upper left')
l = ax3.legend(loc='upper left')

%%writefile fipy_timing.py

"""
Usage: fipy_timing.py [-h] [--N N] [--iterations <iterations>] [--suite <suite>]

Options:
  -h --help                 Show this screen.
  --N=N                     Cell number is N^3 [default: 30]
  --iterations=<iterations> Number of iterations [default: 100]
  --suite=<suite>           Solver suite [default: trilinos]
  
"""

from docopt import docopt
import timeit
import numpy as np
import fipy as fp

arguments = docopt(__doc__, version='Run FiPy timing')
N = int(arguments['--N'])
iterations = int(arguments['--iterations'])
suite = arguments['--suite']

attempts = 3

setup_str = '''
import fipy as fp
import numpy as np
np.random.seed(1)
L = 1.
N = {N:d}
m = fp.GmshGrid3D(nx=N, ny=N, nz=N, dx=L / N, dy=L / N, dz=L / N)
v0 = 1.
for dim in range(3):
    x = np.linspace(0., L, N)
    fx = np.sum(np.array([np.sin(2 * x * np.pi * i * np.random.random()) / i for i in range(1, N)]), axis=0)
    v0 = np.outer(v0, fx).flatten()
v = fp.CellVariable(mesh=m)
v0 = np.resize(v0, len(v)) ## Gmsh doesn't always give us the correct sized grid!
eqn = fp.TransientTerm(1e-3) == fp.DiffusionTerm()
v[:] = v0.copy()

import fipy.solvers.{suite} as solvers
solver = solvers.linearPCGSolver.LinearPCGSolver(precon=None, iterations={iterations}, tolerance=1e-100)

eqn.solve(v, dt=1., solver=solver)
v[:] = v0.copy()
'''

timeit_str = '''
eqn.solve(v, dt=1., solver=solver)
fp.parallelComm.Barrier()
'''

timer = timeit.Timer(timeit_str, setup=setup_str.format(N=N, suite=suite, iterations=iterations))
times = timer.repeat(attempts, 1)

if fp.parallelComm.procID == 0:
    print min(times)

import itertools

runtimes = pd.DataFrame(columns=['iterations', 'N', 'np', 'run time', 'suite'])

Ns = (10, 20, 30, 40, 50, 60, 70, 80)
iterations_ = (100, 200, 400, 800)
nps = (1, 2, 3, 4, 5, 6, 7, 8)
suites = ('trilinos', 'pysparse')
tmpfile = 'tmp.txt'
 
for N, iterations, np_, suite in itertools.product(Ns, iterations_, nps, suites):
    if suite == 'pysparse' and np_ > 1:
        runtime = None
    else:
        !mpirun -np $np_ python fipy_timing.py --N=$N --suite=$suite --iterations=$iterations > $tmpfile
        runtime = float(open(tmpfile).read())
    
    runtimes = runtimes.append({'N' : N, 'iterations' : iterations, 'np' : np_, 'suite' : suite, 'run time' : runtime}, ignore_index=True)
    
runtimes.to_csv('laptop_runtimes_nproc.csv')

df = pd.read_csv('laptop_runtimes_nproc.csv')

f = plt.figure(figsize=(20, 5))
ax1 = f.add_subplot(141)
ax2 = f.add_subplot(142)
ax3 = f.add_subplot(143)
ax4 = f.add_subplot(144)

for iterations, ax in zip((800, 400, 200, 100), (ax1, ax2, ax3, ax4)):
    for N in (10, 20, 30, 40, 50, 60, 70, 80):
        sub_df = df[(df['iterations'] == iterations) & (df['N'] == N) & (df['suite'] == 'trilinos')].sort('np')
        nprocs = sub_df['np']
        speed_up =  np.array(sub_df[sub_df['np'] == 1]['run time']) / np.array(sub_df['run time'])
        ax.semilogx(sub_df['np'], speed_up, label='$N={0}$'.format(N))

    ax.semilogx((1, 8), (1, 8), 'k--')
    ax.set_ylabel("Speed Up")
    ax.set_xlabel("Parallel Processes")
    ax.legend(loc="upper left")
    ax.set_ylim(ymax=3)
    ax.set_title("iterations={0}".format(iterations))