## Import necessary modules and define Lame parameters
using Gridap
using GridapEmbedded
using LinearAlgebra: tr

function lame_parameters(E, ν)
    λ = (E * ν) / ((1 + ν) * (1 - 2 * ν))
    μ = E / (2 * (1 + ν))
    return (λ, ν)
end

## Define material parameters
# Material parameters 1
E1 = 4.0 
ν1 = 0.2
λ1, μ1 = lame_parameters(E1,ν1)
σ1(ε) = λ1*tr(ε)*one(ε) + 2*μ1*ε

# Material parameters 2
E2 = 4.0
ν2 = 0.2
λ2, μ2 = lame_parameters(E2,ν2)
σ2(ε) = λ2*tr(ε)*one(ε) + 2*μ2*ε

# Dirichlet values
g(x,t::Real) = VectorValue(0.0, 0.0, 0.0)*t
g(t::Real) = x -> g(x,t)

b=VectorValue(0,0,-9.81)

## Define model for the background mesh

L = VectorValue(1.0,0.1,0.1)
n = 5
partition = (L[1]/L[3]*n, L[2]/L[3]*n, 10*n)
pmin = Point(0.,0.,0.)
pmax = pmin + L
bgmodel = CartesianDiscreteModel(pmin,pmax,partition)
Ω_bg = Triangulation(bgmodel)

# Identify Dirichlet boundaries
labeling = get_face_labeling(bgmodel)
add_tag_from_tags!(labeling,"support0",[1,3,5,7,13,15,17,19,25])
add_tag_from_tags!(labeling,"support1",[2,4,6,8,14,16,18,20,26])

writevtk(bgmodel, "bgmodel")

##  Define interface geometry

midpoint = 0.51234*L
plane_normal = VectorValue(1,0,0)
geo1 = plane(x0=midpoint, v=plane_normal, name="left")
geo2 = !(geo1, name="right")

## Discretization and integration meshes

cutgeo = cut(bgmodel,union(geo1,geo2))

# Setup interpolation mesh
Ω1_act = Triangulation(cutgeo, ACTIVE, "left")
Ω2_act = Triangulation(cutgeo, ACTIVE, "right")

# Setup integration meshes
Ω1 = Triangulation(cutgeo,PHYSICAL,"left")
Ω2 = Triangulation(cutgeo,PHYSICAL,"right")
Γ = EmbeddedBoundary(cutgeo,"left","right")

# Setup normal vectors
n_Γ = get_normal_vector(Γ)

# Setup Lebesgue measures
order = 1
degree = 2*order
dΩ1 = Measure(Ω1,degree)
dΩ2 = Measure(Ω2,degree)
dΓ = Measure(Γ,degree)


## Setup FESpace
V1 = TestFESpace(Ω1_act,
                 ReferenceFE(lagrangian,VectorValue{3,Float64},order),
                 conformity=:H1,
                 dirichlet_tags=["support0"])
V2 = TestFESpace(Ω2_act,
                 ReferenceFE(lagrangian,VectorValue{3,Float64},order),
                 conformity=:H1)

U1 = TransientTrialFESpace(V1,g)
U2 = TransientTrialFESpace(V2)

V = MultiFieldFESpace([V1,V2])
U = TransientMultiFieldFESpace([U1,U2])

## Set up weak variation formulation

# Setup stabilization parameters
meas_K1 = get_cell_measure(Ω1, Ω_bg)
meas_K2 = get_cell_measure(Ω2, Ω_bg)
meas_KΓ = get_cell_measure(Γ, Ω_bg)

γ_hat = 2
κ1 = CellField( (E2*meas_K1) ./ (E2*meas_K1 .+ E1*meas_K2), Ω_bg)
κ2 = CellField( (E1*meas_K2) ./ (E2*meas_K1 .+ E1*meas_K2), Ω_bg)
β  = CellField( (γ_hat*meas_KΓ) ./ ( meas_K1/E1 .+ meas_K2/E2 ), Ω_bg)

# Jump and mean operators for this formulation
jump_u(u1,u2) = u1 - u2
mean_t(u1,u2) = κ1*(σ1∘ε(u1)) + κ2*(σ2∘ε(u2))

m(t,(u1,u2),(v1,v2))=∫( u1 ⊙ v1 ) * dΩ1 + ∫( u2 ⊙ v2 ) * dΩ2

# Weak form
a(t,(u1,u2),(v1,v2)) =
  ∫( ε(v1) ⊙ (σ1∘ε(u1)) ) * dΩ1 +
  ∫( ε(v2) ⊙ (σ2∘ε(u2)) ) * dΩ2 + 
  ∫( β*jump_u(v1,v2)⋅jump_u(u1,u2)
     - n_Γ⋅mean_t(u1,u2)⋅jump_u(v1,v2)
     - n_Γ⋅mean_t(v1,v2)⋅jump_u(u1,u2) )* dΓ

l(t,(v1,v2)) = ∫(b ⊙ v1) * dΩ1 + ∫(b ⊙ v2) * dΩ2 

U10 = U1(0.0)
U20 = U2(0.0)
U0 = U(0.0)
uh10 = interpolate_everywhere(VectorValue(0,0,0),U10)
uh20 = interpolate_everywhere(VectorValue(0,0,0),U20)
uh0 = interpolate_everywhere([uh10,uh20],U0)

# FE problem
op = TransientAffineFEOperator(m,a,l,U,V)

linear_solver = LUSolver()

Δt = 0.05
θ = 0.5
ode_solver = ThetaMethod(linear_solver,Δt,θ)

t₀ = 0.0
T = 10.0
uht = solve(ode_solver,op,uh0,t₀,T)

##