from firedrake import *
import numpy as np

# files to save data
save_phi_s = File("data test/phi_s.pvd")
save_phi_3s = File("data test/phi_3s.pvd")


# surface mesh (sigma = H0)
surface_mesh = RectangleMesh(10,10,5.0,2.0)
# 3D mesh
mesh = ExtrudedMesh(surface_mesh, layers=10, layer_height=0.05, extrusion_type='uniform')

# Function space for the 2D surface mesh
Vsurf = FunctionSpace(surface_mesh,"CG", 1)
#Function space for the 3D mesh
V = FunctionSpace(mesh, "CG", 1)
# NOTE: if we use insted CG2, then the value is not exactly the same.

# phi_3s : the 3D function that we want to be homogeneous in sigma
phi_3s=Function(V)
# phi_s: 3D function which is updated only at the surface. We will use its value at the surface (and bottom) as boundary conditions for phi_3D.
phi_s = Function(V)

# test function
v=TestFunction(V)

# unknown
phi = TrialFunction(V)

# normal
n=FacetNormal(mesh)

# We initialise phi_2D. It is initialise in the whole domain but it does not matter: only the surface and bottom values will be used.
phi_s.interpolate(Expression("x[0]"))
# phi_3D is initialy set to be zero everywhere.
phi_3s.assign(0.0)

# Save the initial conditions
save_phi_3s << phi_3s
save_phi_s << phi_s

# Define the boundary condition : in the function space V, we define the top and bottom surfaces to be equal to phi_s at these surfaces.
bc_top = DirichletBC(V, phi_s, 'top')
bc_bottom = DirichletBC(V, phi_s, 'bottom')


# Weak formulation : see my comments in the email
a =(dot(grad(phi),grad(v)))*dx
b=(v.dx(0)*phi_3s.dx(0) + v.dx(1)*phi_3s.dx(1))*dx

# solve the WF with boundary conditions defined above
solve(a==b, phi_3s, bcs=[bc_bottom,bc_top])

# save the updated values : phi_3s is indeed homogeneous in the vertical.
save_phi_3s << phi_3s
save_phi_s << phi_s