# Copyright (C) 2008 Ilmar Wilbers
#
# This file is part of DOLFIN.
#
# DOLFIN is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# DOLFIN is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with DOLFIN. If not, see <http://www.gnu.org/licenses/>.
#
# Modified by Anders Logg, 2008.
#
# First added:  2008-06-04
# Last changed: 2008-07-21

# Be careful to first run the benchmark once in order for
# the Python extensions to be created, or else the timings
# will not be correct.

from dolfin import *
from time import time
import sys

def make_form(name, mesh):
    globals()['mesh'] = mesh
    execfile("../forms/" + name + '.form', globals())
    try:
        return a
    except:
        print "No object 'a' to return in file %s.form" %name
        return None

def bench_form(form, mesh, reps=1):
    totaltime = 0.0
    t0 = time()
    A = assemble(form, mesh)
    totaltime += time() - t0
    for i in range(reps - 1):
        t0 = time()
        assemble(form, mesh, tensor=A, reset_sparsity=False)
        totaltime += time() - t0
    return totaltime / float(reps)

def make_mesh(name, dim):
    if dim == 3:
        N = 32
        mesh = UnitCube(N, N, N)
        return mesh
    else:
        N = 256
        mesh = UnitSquare(N, N)
        return mesh

if __name__ == "__main__":
    try:
        reps = int(sys.argv[1])
    except:
        print 'Usage: %s number_of_repetitions [1]' %sys.argv[0]
        reps = 1

# Backends
backends = ["uBLAS", "PETSc", "Tpetra", "Eigen"]

# Forms
forms = ["Elasticity3D",
         "PoissonP1",
         "PoissonP2",
         "PoissonP3",
         "THStokes2D",
         "NSEMomentum3D",
         "StabStokes2D"]

dolfin_set("output destination", "silent")
results = Table("Assembly benchmark")

for backend in backends:
    dolfin_set("linear algebra backend", backend)
    for form in forms:
        dim = 2 if not form.find("3D") > -1 else 3
        m = make_mesh(form, dim)
        a = make_form(form, m)
        print "Assembling %s with %s" % (form, backend)
        t = bench_form(a, m, reps=reps)
        results.set(backend, form, t)

dolfin_set("output destination", "terminal")
print ""

results.disp()
list_timings(TimingClear_keep,
             [TimingType_wall, TimingType_user, TimingType_system])
