"""
This is a simple, but inefficient implementation of the nonlinear
shallow water equations, applied to the Williamson 2 test case
"""

# Copyright (C) 2013 Andrew McRae

from dolfin import *

f = Expression("2*7.292e-5*x[2]/6.37122e6") # Coriolis frequency (1/s)
g = Constant(9.80616)   # gravitational constant (m/s^2)
dt = 3600.0     # timestep (s)
T = 432000.0    # total time (s)
theta = 0.5     # 0 = Backward Euler, 1 = Forward Euler

mesh = Mesh("meshes/sphere_ico3_earth.xml.gz")
global_normal = Expression(("x[0]", "x[1]", "x[2]"))
mesh.init_cell_orientations(global_normal)

class ElementNormal(Expression):
    def __init__(self, mesh):
        self.mesh = mesh
    def eval_cell(self, values, x, ufc_cell):
        orientation = mesh.cell_orientations()[ufc_cell.index]
        cell = Cell(self.mesh, ufc_cell.index)
        n = cell.cell_normal()
        if (orientation == 0):
            values[0] = n[0]
            values[1] = n[1]
            values[2] = n[2]
        else:
            values[0] = -n[0]
            values[1] = -n[1]
            values[2] = -n[2]
    def value_shape(self):
        return (3,)
en_exp = ElementNormal(mesh)

# set up function spaces
E = FunctionSpace(mesh, 'CG', 2)
V = FunctionSpace(mesh, 'BDM', 1)
Q = FunctionSpace(mesh, 'DG', 0)
VQEV = MixedFunctionSpace([V, Q, E, V])
PlotSpace = VectorFunctionSpace(mesh, 'CG', 1)

# set up velocity, surface height and topography fields
uexpr = Expression(("38.61068276698372*-x[1]/6.37122e6", "38.61068276698372*x[0]/6.37122e6", "0.0"))
Dexpr = Expression("(2.94e4 - (6.37122e6 * 7.292e-5 * 38.61068276698372 + pow(38.61068276698372, 2)/2.0)*(x[2]*x[2]/(6.37122e6*6.37122e6)))/9.80616")
bexpr = Expression("0.0")
u = project(uexpr, V)
D = project(Dexpr, Q)
b = project(bexpr, Q)

def gradperp(psi):
    return cross(en_exp, grad(psi))
def myperp(u):
    return cross(en_exp, u)

# Generate guess for initial Newton iteration - du = 0, dh = 0, q = q_0, F = F_0
F = TrialFunction(V)
w = TestFunction(V)
F_0 = Function(V)
solve(dot(w, F)*dx==dot(w, D*u)*dx, F_0) # Generate F_0

q = TrialFunction(E)
gamma = TestFunction(E)
a = gamma*q*D*dx
L = gamma*f*dx - dot(gradperp(gamma), u)*dx
q_0 = Function(E)
solve (a==L, q_0) # Generate q_0

# Generate combined field in V x Q x E x V, initialised as above
(du0, dD0, q0, F0) = TrialFunctions(VQEV)
(dut, dDt, qt, Ft) = TestFunctions(VQEV)
a = (dot(du0, dut) + dD0*dDt + q0*qt + dot(F0, Ft))*dx
L = (q_0*qt + dot(F_0, Ft))*dx
dudDqF = Function(VQEV)
solve(a==L, dudDqF)

uplot = project(u, PlotSpace)
plot(uplot)
plot(D)
t = 0.0
while t < (T-dt/2):
    # Fully implicit timestepping in terms of du, dh, q, F
    (du, dD, q, F) = split(dudDqF)  # "new" variables
    (w, phi, gamma, w2) = TestFunctions(VQEV)
    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
    
    F1 = (dot(w, du) - dt*div(w)*(g*Db_i + 0.5*dot(u_i, u_i))
                + dt*(q-(dt/2)*dot(u_i, grad(q)))*dot(w, myperp(F)))*dx
    F2 = (phi*dD + dt*phi*div(F))*dx
    F3 = (gamma*q*D_i + dot(gradperp(gamma), u_i) - gamma*f)*dx
    F4 = (dot(w2, F) - dot(w2, D_i*u_i))*dx
    solve(F1+F2+F3+F4 == 0, dudDqF)
    (du, dD, q, F) = dudDqF.split(deepcopy=True)
    u.vector().axpy(1, du.vector())
    D.vector().axpy(1, dD.vector())
    print "t =", t
    t += dt
    uplot.assign(project(u, PlotSpace))
    plot(uplot)
    plot(D)
