"""
Linear shallow water code for mimetic element pairs.

The equations are the linear shallow water equations.  Let u be
velocity, D be depth _perturbation_, g = gravity, H = base layer depth

du/dt + fu^\perp = -g grad(D)
dD/dt + H div(u) = 0

In the continuous case, if we pick a velocity field u = curl(psi),
then div(u) = 0, so dD/dt is initially 0.  If you then choose D to
satisfy fu^\perp = -g grad(D), du/dt is initially 0 also.  Then both u
and D remain fixed for all time.  This is known as "geostrophic
balance".  One of the special properties of the spatial discretisation
is that it represents this balance.

With implicit midpoint timestepping, the scheme conserves energy -
this is verified in.
"""

# Copyright (C) 2013 Andrew McRae and Imperial College London
# Written by Andrew McRae and Colin J. Cotter
# Modified by Marie E. Rognes (meg@simula.no), 2013

from dolfin import *

parameters["form_compiler"]["cpp_optimize"] = True
parameters["form_compiler"]["optimize"] = True

outward_normals_code = """
class OutwardNormal : public Expression
{
public:

boost::shared_ptr<Mesh> mesh;

OutwardNormal() : Expression(3)
{
}

void eval(Array<double>& values, const Array<double>& x,
          const ufc::cell& c) const
  {
   assert(mesh);
   const Cell cell(*mesh, c.index);

   // Extract local normal and orientation (already computed)
   const Point& local_normal(cell.cell_normal());
   const int orientation = mesh->cell_orientations()[c.index];

   // Set values (1 means flipped, so then we flip back)
   if (orientation == 0)
   {
     values[0] = local_normal.x();
     values[1] = local_normal.y();
     values[2] = local_normal.z();
   } else
   {
     values[0] = - local_normal.x();
     values[1] = - local_normal.y();
     values[2] = - local_normal.z();
   }
  }
};"""

def initialize_parameters(verbose=True):
    # Create application parameters
    params = Parameters("Discretization")
    params.add("family", "RT")
    params.add("order", 1)
    params.add("meshid", "sphere_ico5")
    params.add("store_pvds", True)
    params.add("T", 10.0)
    params.add("dt", 0.05)
    params.add("directory", "default")

    # Read parameters from the command-line
    params.parse()
    if params["directory"] == "default":
        default = "%s%d_%s" % (params["family"], params["order"],
                               params["meshid"])
        params["directory"] = default

    info(params, verbose)

    # Return parameters
    return params

def create_function_spaces(mesh, family, order):
    S = FunctionSpace(mesh, family, order)
    V = FunctionSpace(mesh, "DG", order-1)
    PlottingMesh = VectorFunctionSpace(mesh, "CG", order)
    return (S, V, PlottingMesh)

def myperp(u):
    return cross(en_exp, u)

def initial_conditions(S, V):
    u0 = interpolate(Constant((0.0,0.0,0.0)), S)
    D0 = project(Expression("exp(-pow((-x[1]-1.0)/0.1, 2))"), V)
    return (u0, D0)

def energy(u, D, H, g):
    kinetic = 0.5*H*dot(u, u)*dx
    potential = 0.5*g*D*D*dx
    return (kinetic, potential)

def output_energy(energyfile, t, kinetic, potential):
    energyfile.write("times.append(%g);\n" % t)
    energyfile.write("kinetics.append(%g);\n" % kinetic)
    energyfile.write("potentials.append(%g);\n" % potential)

def main():

    # Setup parameters
    params = initialize_parameters()

    # Store parameters
    storage = "results/%s" % params["directory"]
    paramsfile = File("%s/application_parameters.xml" % storage)
    paramsfile << params

    # Set-up storage for energies
    energyfilename = "results/%s/energies.py" % params["directory"]
    info_blue("Storing energies to %s" % energyfilename)
    energyfile = open(energyfilename, "w")
    energyfile.write("# Results from linear shallow water test case.\n")
    energyfile.write("kinetics = [];\n")
    energyfile.write("potentials = [];\n")
    energyfile.write("times = [];\n")

    # Initialize mesh
    mesh = Mesh("../mixed-poisson/hdiv-l2/meshes/%s.xml.gz" % params["meshid"])
    global_normal = Expression(("x[0]", "x[1]", "x[2]"))
    mesh.init_cell_orientations(global_normal)

    # Set-up function spaces
    (S, V, PlottingMesh) = create_function_spaces(mesh, params["family"], params["order"])
    W = S * V

    # Setting up perp operator
    outward_normals = Expression(outward_normals_code)
    outward_normals.mesh = mesh
    perp = lambda u: cross(outward_normals, u)

    # Extract initial conditions
    (u_, D_) = initial_conditions(S, V)

    # Extract some parameters for discretization
    dt = Constant(params["dt"])
    f = Expression("x[2]")
    H = 1.0
    g = 1.0

    # Implicit midoint scheme discretization in time
    (u, D) = TrialFunctions(W)
    (w, phi) = TestFunctions(W)
    u_mid = 0.5*(u + u_)
    D_mid = 0.5*(D + D_)
    F = (inner(u - u_, w) - dt*div(w)*g*D_mid + dt*f*inner(perp(u_mid), w)
         + inner(D - D_, phi) + dt*H*div(u_mid)*phi)*dx
    (a, L) = system(F)

    # Preassemble matrix (because we can)
    A = assemble(a)

    # Define energy functional
    (kinetic_func, potential_func) = energy(u_, D_, H, g)

    # Setup solution function for current time
    uD = Function(W)

    # Predefine b (for the sake of reuse of memory)
    b = Vector(W.dim())

    # Set-up linear solver (so that we can reuse LU)
    solver = LUSolver(A)
    solver.parameters["same_nonzero_pattern"] = True
    solver.parameters["reuse_factorization"] = True

    # Set-up some pvd storage
    if params["store_pvds"]:
        ufile = File("%s/pvds/u.pvd" % storage)
        Dfile = File("%s/pvds/D.pvd" % storage)

    # Set-up some time related variables
    k = 0
    t = 0.0
    T = params["T"]
    dt = float(dt)

    # Output initial energy
    E_k = assemble(kinetic_func)
    E_p = assemble(potential_func)
    print t, E_k, E_p, E_k + E_p, D_.vector().min(), D_.vector().max()
    output_energy(energyfile, t, E_k, E_p)

    # Time loop
    while(t < (T - 0.5*dt)):

        # Assemble right-hand side, reuse b
        assemble(L, tensor=b)

        # Solve system
        solver.solve(uD.vector(), b)

        # Update previous solution
        u, D = uD.split()
        u_.assign(u)
        D_.assign(D)

        # Update time and counter
        t += dt
        k += 1

        # Output current energy and max/min depth
        E_k = assemble(kinetic_func)
        E_p = assemble(potential_func)
        print t, E_k, E_p, E_k + E_p, D_.vector().min(), D_.vector().max()
        output_energy(energyfile, t, E_k, E_p)

        # Store solutions to xml and pvd
        solutionfile = File("%s/xmls/uD_%d.xml.gz" % (storage, k))
        solutionfile << uD

        if params["store_pvds"]:
            uplot = project(u,PlottingMesh)
            ufile << uplot
            Dfile << D

if __name__ == "__main__":

    tic()
    main()
    print toc()
