#
# Solve shallow water equations (SWE) with a buoy.
#
# Last update: 02.12.15
#
from firedrake import *
from parameters import *
from mesh import *
from IC_etaR import *
from solvers import *
from energy import *
from heavyside import *
import time

#op2.init(log_level='WARNING')		# Set the log level. Options: DEBUG, INFO, WARNING, ERROR, CRITICAL
#parameters["assembly_cache"]["enabled"] = False

# Parameter values
(sluice_gate, wave_maker, solvers_print) = get_cases();
(L, H0, g, rho, Mass, a, Lp, Nx, Np, dt, T) = parameter_values();
if sluice_gate == 'True': (x1, x2, h1, Ts) = etaR_input(H0);
if wave_maker == 'True': (Ampl, k, omega) = dispersion_relation(L, H0, g, Nx);

# Create mesh 
(mesh, coords) = create_mesh(Nx, L);

# Define functions
V = FunctionSpace(mesh, "CG", 1)
V_DG0 = FunctionSpace(mesh, "DG", 0)

eta0 = Function(V)
phi0 = Function(V)
eta1   = Function(V)
phi1 = Function(V)
phi0_5 = Function(V)
mu0_5 = Function(V)
lambda0_5 = Function(V)
h = Function(V)

Hb = Function(V)
etaR = Function(V)
dR_dt = Function(V)

step_b = Function(V_DG0)		# Heavyside step function: 0 in water part, 1 in buoy part.
step_w = Function(V_DG0)		# Heavyside step function: 1 in water part, 0 in buoy part.

Z0   = Constant(0.0)
W0   = Constant(0.0)
Z1   = Constant(0.0)
W1   = Constant(0.0)
W0_5 = Constant(0.0)
xp   = Constant(Lp)

eta = TrialFunction(V)
phi = TrialFunction(V)
mu = TrialFunction(V)
v = TestFunction(V)

# Buoy's hull shape
(Hb, H, Zbar) = def_Hb(Hb, coords.dat.data, L, Lp, H0, rho, Mass, a, Nx);

# Heavyside step function
step_b = def_step(step_b, coords.dat.data, Lp, 'True');
step_w = def_step(step_w, coords.dat.data, Lp, 'False');

# Time iteration
t = 0
if sluice_gate == 'True': set_E0 = 'True'
else: set_E0 = 'False'

# Initial condition for sluice gate
if sluice_gate == 'True':
    etaR_expr = def_etaR(coords.dat.data, H0, h1, x1, x2, Ts, t);
    etaR.interpolate(etaR_expr)
    eta0 = eta0_eq_etaR(eta,eta0,etaR,v);
elif wave_maker == 'True':
    dR_dt_expr = Expression("A*sin(w*t)", A=Ampl, w=omega, t=t)

# Weak formulation
(mu_solver0_5, mu0_5, Mb_Minv_A) = solver_mu(eta0, phi0, Z0, W0, mu0_5, mu, v, step_b, dt, Hb, g, rho, Mass);   # Linear system for Lagrange multiplier mu
phi_solver0_5 = solver_phi(phi0_5, phi0_5, phi0, eta0, mu0_5, step_b, etaR, phi, v, 0.5*dt, g, solvers_print);  # Give phi0_5 instead of phi1
eta_solver1 = solver_eta(eta1, eta1, eta0, phi0_5, eta, v, dt, Hb, H0, dR_dt, solvers_print);	                # Give eta1 instead of eta0_5
phi_solver1 = solver_phi(phi1, phi0_5, phi0_5, eta1, mu0_5, step_b, etaR, phi, v, 0.5*dt, g, solvers_print);    # Give eta1 instead of eta0 and phi0_5 instead of phi0

lambda0_5.assign((2/dt)*mu0_5)
h.assign(Hb+eta0)

# Write data to files
phi_file = File("phi.pvd")
eta_file   = File("eta.pvd")
lambda_file = File("lambda.pvd")
h_file = File("h.pvd")

phi_file << phi0
eta_file << eta0
lambda_file << lambda0_5
h_file << h

xp_file = open("waterline.csv", "w")
xp_file.write('%-2s, %-2s\n' % ('xp','hp'))
xp_file.write('%-10s, %-10s\n' % (str(float(xp)),str(H0)))

# Initial energy
E0_w = 0.5*rho*assemble( (Hb*abs(grad(phi0))**2 + g*eta0**2)*dx )   # Energy of the water
E0_b = 0.5*Mass*float(W0)**2										# Energy of the buoy
E0   = E0_w + E0_b				                                    # Total energy: water + buoy

E_file = open("energy.txt", "w")
E = find_energy(phi0, eta0, W0, Z0, Hb, rho, g, Mass, L, t, E0, E_file);

while(t < T-dt):
    print t, E

    # Time update of sluice gate etaR and wave maker R(t) = A/w*(1-cos(w*t))
    etaR.interpolate(Expression("0"))
    dR_dt.interpolate(Expression("0"))

    if (sluice_gate == 'True' and t < Ts):
        etaR_expr.t = t+0.5*dt
        etaR.interpolate(etaR_expr)
    elif wave_maker == 'True':
        dR_dt_expr.t = t+0.5*dt
        dR_dt.interpolate(dR_dt_expr)

    t += dt

    # Update RHS
    rhs = find_rhs(Mb_Minv_A, dt, g, eta0, phi0, v, step_b, Z0, W0);

	# Solve the variational problem
    (phi0, eta0, W0, Z0) = solvers_SV(rhs, mu0_5, mu_solver0_5, phi_solver0_5, eta_solver1, phi_solver1, eta0, phi0, Z0, W0, eta1, phi1, Z1, W1, W0_5, step_b, rho, Mass, dt);

    lambda0_5.assign((2/dt)*mu0_5)
    h.assign(Hb+eta0)

    # Calculate waterline point
    xp.assign(Lp + (Z0 - eta0.dat.data[Np])/tan(a))
    xp_file.write('%-10s, %-10s\n' % (str(float(xp)),str(H0+eta0.dat.data[Np])))

	# Monitor energy in time
    if sluice_gate == 'True':
        if (t > Ts and set_E0 == 'True'):    # As soon the sluice gate is released (i.e. t>Ts), calculate E0 again
    	    E0 = 0.5*rho*assemble( (Hb*abs(grad(phi0))**2 + g*eta0**2)*dx ) + 0.5*Mass*float(W0)**2
            set_E0 = 'False'

    E = find_energy(phi0, eta0, W0, Z0, Hb, rho, g, Mass, L, t, E0, E_file);

	# Write data to files
    eta_file << eta0
    phi_file << phi0
    lambda_file << lambda0_5
    h_file << h

E_file.close()
xp_file.close()