using Gridap
using ForwardDiff
using GridapODEs.ODETools
using GridapODEs.TransientFETools

import Gridap: ∇
import GridapODEs.TransientFETools: ∂t

θ = 0.5

u(x,t) = (1.0-x[1])*x[1]*(1.0-x[2])*x[2]*t
u(t) = x -> u(x,t)

f(t) = x -> ∂t(u)(x,t)-Δ(u(t))(x) # or ∂t(u)(t)(x)-Δ(u(t))(x)

domain = (0,1,0,1)
partition = (4,4)
model = CartesianDiscreteModel(domain,partition)

order = 2

reffe = ReferenceFE(lagrangian,Float64,order)
V0 = FESpace(
  model,
  reffe,
  conformity=:H1, 
  dirichlet_tags="boundary"
)

U = TransientTrialFESpace(V0,u)

Ω = Triangulation(model)
degree = 2*order
dΩ = Measure(Ω,degree)

a(u,v) = ∫( ∇(v)⋅∇(u) )dΩ
b(v,t) = ∫( v*f(t) )dΩ
m(u,v) = ∫( v*u )dΩ

res(t,u,ut,v) = a(u,v) + m(ut,v) - b(v,t)
jac(t,u,ut,du,v) = a(du,v)
jac_t(t,u,ut,dut,v) = m(dut,v)

op = TransientFEOperator(res,jac,jac_t,U,V0)

t0 = 0.0
tF = 1.0
dt = 0.1

U0 = U(0.0)
uh0 = interpolate_everywhere(u(0.0),U0)

ls = LUSolver()
odes = ThetaMethod(ls,dt,θ)
solver = TransientFESolver(odes)
sol_t = solve(solver,op,uh0,t0,tF)

for (uh_tn, tn) in sol_t
  # Here we have the solution uh_tn at tn
end