"""
An efficient implementation of the nonlinear shallow water equations,
applied to the Williamson 5 test case.
"""

# Copyright (C) 2013 Andrew McRae, with code borrowed from Rognes/Cotter

from dolfin import *
import numpy

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("T", 1296000.0)     # seconds
    params.add("dt", 600.0)
    params.add("maxiter", 5)    # number of Newton iterations
    params.add("theta", 0.5)    # 0 = Backward Euler, 1 = Forward Euler

    # Read parameters from the command-line
    params.parse()
    info(params, verbose)

    # Return parameters
    return params

def create_function_spaces(mesh, family, order):
    if family == "BDM":
        E = FunctionSpace(mesh, "CG", order+1)
    elif family == "RT":
        E = FunctionSpace(mesh, "CG", order)
    else:
        pass
    V = FunctionSpace(mesh, family, order)
    Q = FunctionSpace(mesh, "DG", order-1)
    PlotSpace = VectorFunctionSpace(mesh, "CG", order)
    return (E, V, Q, PlotSpace)

def initial_conditions(V, Q):
    uexpr = Expression(("-20.0*x[1]/6.37122e6", "20.0*x[0]/6.37122e6", "0.0"))
    Dexpr = Expression("5960 - ((6.37122e6 * 7.292e-5 * 20.0 + pow(20.0,2)/2.0)*(x[2]*x[2]/(6.37122e6*6.37122e6)))/9.80616 - (2000 * (1 - sqrt(min(pow(pi/9.0,2),pow(atan2(x[1]/6.37122e6,x[0]/6.37122e6)+1.0*pi/2.0,2)+pow(asin(x[2]/6.37122e6)-pi/6.0,2)))/(pi/9.0)))")
    bexpr = Expression("2000 * (1 - sqrt(min(pow(pi/9.0,2),pow(atan2(x[1]/6.37122e6,x[0]/6.37122e6)+1.0*pi/2.0,2)+pow(asin(x[2]/6.37122e6)-pi/6.0,2)))/(pi/9.0))")
    u0 = project(uexpr, V)
    D0 = project(Dexpr, Q)
    b0 = project(bexpr, Q)
    return (u0, D0, b0)

def main():

    # Setup parameters
    params = initialize_parameters()
    
    # Parameters
    f = Expression("2*7.292e-5*x[2]/6.37122e6") # Coriolis frequency (1/s)
    g = Constant(9.80616)   # gravitational constant (m/s^2)
    T = params["T"]
    dt = Constant(params["dt"])
    theta = Constant(params["theta"])
    maxiter = params["maxiter"]
    
    # Define global normal
    global_normal = Expression(("x[0]", "x[1]", "x[2]"))

    # Define meshes
    meshes = ["sphere_ico3_earth", "sphere_ico4_earth", "sphere_ico5_earth", "sphere_ico6_earth"]

    for (family, order) in [("RT", 1), ("BDM", 1)]:
        for (i,meshid) in enumerate(meshes):
            meshname = "../williamson2/meshes/%s.xml.gz" % meshid
            mesh = Mesh(meshname)
            mesh.init_cell_orientations(global_normal)

            # Define function spaces and basis functions
            (E, V, Q, PlotSpace) = create_function_spaces(mesh, family, order)
            VQ = MixedFunctionSpace([V, Q])
            plotu = Function(PlotSpace)

            # Set up perp operators
            outward_normals = Expression(outward_normals_code)
            outward_normals.mesh = mesh
            myperp = lambda u: cross(outward_normals, u)
            gradperp = lambda psi: cross(outward_normals, grad(psi))

            # set up velocity, surface height and topography fields
            (u, D, b) = initial_conditions(V, Q)
            # keep initial data to compare against at the end
            u0 = Function(V)
            D0 = Function(Q)
            u0.assign(u)
            D0.assign(D)

            # get average height
            temp = interpolate(Constant(1.0), Q)
            H = assemble(D*dx)/assemble(temp*dx)

            t = 0.0

            du = Function(V)
            dD = Function(Q)
            u_i = u + (1-theta)*du      # implicit velocity
            D_i = D + (1-theta)*dD      # implicit fluid thickness
            Db_i = D + b + (1-theta)*dD # implicit surface height
            
            # Forms for generating 'best guess' for F and q
            F_ = TrialFunction(V)
            w_ = TestFunction(V)
            a_8 = dot(w_, F_)*dx
            L_8 = dot(w_, D_i*u_i)*dx
            # Preassemble matrix (because we can)
            A_8 = assemble(a_8)
            # Set up solution function
            F = Function(V)
            # Predefine b (for the sake of reuse of memory)
            b_8 = Vector(V.dim())
            # Set-up linear solver (so that we can reuse LU)
            solver_8 = LUSolver(A_8)
            solver_8.parameters["same_nonzero_pattern"] = True
            solver_8.parameters["reuse_factorization"] = True

            q_ = TrialFunction(E)
            gamma = TestFunction(E)
            a_9 = gamma*q_*D_i*dx
            L_9 = gamma*f*dx - dot(gradperp(gamma), u_i)*dx
            # Set up solution function        
            q = Function(E)
            
            (deltadu, deltadD) = TrialFunctions(VQ)
            (w, phi) = TestFunctions(VQ)
            a1 = (dot(w,deltadu) + (1-theta)*dt*f*dot(w,myperp(deltadu))
                   - (1-theta)*dt*g*div(w)*deltadD)*dx
            a2 = (deltadD*phi + (1-theta)*dt*H*phi*div(deltadu))*dx
            L1 = (-dot(w,du) - dt*(q-(dt/2)*dot(u,grad(q)))*dot(w,myperp(F))
               + dt*div(w)*(g*Db_i + 0.5*dot(u_i,u_i)))*dx
            L2 = (-phi*dD - dt*phi*div(F))*dx
            a = a1 + a2
            L = L1 + L2
            # Preassemble matrix (because we can)
            A = assemble(a)
            # Set up solution function for current time
            deltaduD = Function(VQ)
            # Predefine b (for the sake of reuse of memory)
            b1 = Vector(VQ.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

            dt = float(dt)
            while t < (T-dt/2):
                # Implicit timestepping with a manual Newton solve, Jacobian
                # corresponding to the LSW equations
                du.vector().zero()
                dD.vector().zero()
                niter = 0
                while (niter < maxiter):
                    # Generate 'best guess' for F and q
                    assemble(L_8, tensor=b_8)
                    solver_8.solve(F.vector(), b_8)

                    solve(a_9 == L_9, q)

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

                    # Solve system
                    solver.solve(deltaduD.vector(), b1)
                    
                    # Not strictly necessary, but very useful...
                    b2 = assemble(L2)
                    res = b2.norm("l2")
                    print "Depth residual:", res
                    
                    deltadu, deltadD = deltaduD.split(deepcopy=True)
                    du.vector().axpy(1, deltadu.vector())
                    dD.vector().axpy(1, deltadD.vector())
                    niter += 1

                u.vector().axpy(1, du.vector())
                D.vector().axpy(1, dD.vector())
                t += dt
                print "t =", t
                #plotu.assign(project(u,PlotSpace))
                #plot(plotu)
                #plot(D)
            # Writing output files
            File("results/u"+family+str(order)+"_"+str(i)+".xml") << u
            # add terrain back on to layer depth before outputting
            D.vector().axpy(1, b.vector())
            File("results/D"+family+str(order)+"_"+str(i)+".xml") << D
            File("results/D"+family+str(order)+"_"+str(i)+".pvd") << D

if __name__ == "__main__":

    tic()
    main()
    print toc()
