from dolfinx import mesh, fem, default_scalar_type, io
import ufl
import basix
from dolfinx.cpp.nls.petsc import NewtonSolver
from mpi4py import MPI
from petsc4py import PETSc
import numpy as np
from dolfinx_materials.quadrature_map import QuadratureMap
from dolfinx_materials.solvers import NonlinearMaterialProblem
from dolfinx_materials.material.mfront import MFrontMaterial
from dolfinx_materials.utils import symmetric_tensor_to_vector
import os
current_path = os.getcwd()


L, H = 0.1, 0.1
Nx, Ny = 20, 20

domain = mesh.create_rectangle(MPI.COMM_WORLD, [[0,0], [L,H]], [Nx,Ny], mesh.CellType.quadrilateral)
x = ufl.SpatialCoordinate(domain)
def axi_grad(u):
    return ufl.as_vector([u[0].dx(0), u[1].dx(1), u[0]/x[0], (u[0].dx(1)+u[1].dx(0))/2])

mat_prop = {
    "YoungModulus": 6e9,
    "PoissonRatio": 0.2,
    "BiotCoefficient": 0.6,
    "BiotN": 3.1e10,
    "LiquidPermeability": 5e-16,
    "InitialPorosity": 0.2,
}
material = MFrontMaterial(
    os.path.join(current_path, "src/libBehaviour.so"),
    "PoroElasticity_Maxime",
    hypothesis="plane_strain",
    material_properties = mat_prop,
)

disp_el = basix.ufl.element("Lagrange", domain.basix_cell(), 1, shape=(2,))
fluid_el = basix.ufl.element("Lagrange", domain.basix_cell(), 1)
V = fem.functionspace(domain, basix.ufl.mixed_element([disp_el, fluid_el]))
v = fem.Function(V, name="Unknown")

def bottom(x):
    return np.isclose(x[1], 0)
def top(x):
    return np.isclose(x[1], H)
def left(x):
    return np.isclose(x[0], 0)


bottom_facets = mesh.locate_entities_boundary(domain, 1, bottom)
bottom_dofs_y = fem.locate_dofs_topological(V.sub(0).sub(1), 1, bottom_facets)
top_facets = mesh.locate_entities_boundary(domain, 1, top)
pf_ini = 0e6
top_dofs_pf = fem.locate_dofs_topological(V.sub(1), 1, top_facets)
bc_pf_top = fem.dirichletbc(pf_ini, top_dofs_pf, V.sub(1))

bcs = [
    fem.dirichletbc(default_scalar_type(0), bottom_dofs_y, V.sub(0).sub(1)),
    bc_pf_top
]

facet_tags = mesh.meshtags(domain, 1, top_facets, np.full_like(top_facets, 1))

ds = ufl.Measure("ds", domain=domain, subdomain_data=facet_tags, metadata={"quadrature degree": 2})
f = fem.Constant(domain, np.zeros((2,)))

qmap = QuadratureMap(domain, 2, material)

(u, pl) = ufl.split(v)
qmap.register_gradient("Strain", symmetric_tensor_to_vector(ufl.sym(ufl.grad(u))))
qmap.register_external_state_variable("LiquidPressure", pl)
qmap.register_gradient("LiquidPressureGradient", ufl.grad(pl))

sig = qmap.fluxes["Stress"]
w = qmap.fluxes["FluidFlux"]
ml = qmap.internal_state_variables["LiquidMass"]
ml_n = ml.copy()
dt = fem.Constant(domain, 0.0)

v_ = ufl.TestFunction(V)
(u_, pl_) = ufl.split(v_)
dv = ufl.TrialFunction(V)

F = (ufl.dot(sig, symmetric_tensor_to_vector(ufl.sym(ufl.grad(u_))))) * qmap.dx + ( (ml-ml_n)/dt*pl_ - ufl.dot(w, ufl.grad(pl_)))*qmap.dx -ufl.dot(f, u_)*ds(1)

Jac = qmap.derivative(F, v, dv)

problem = NonlinearMaterialProblem(qmap, F, Jac, v, bcs)

qmap.update()
ml.x.petsc_vec.copy(ml_n.x.petsc_vec)

newton = NewtonSolver(MPI.COMM_WORLD)
newton.rtol = 1e-6
newton.atol = 1e-6
newton.convergence_criterion = "incremental"
newton.report = True
newton.max_it = 50

Nincr = 20
load_steps = np.linspace(0.0, 1.0, Nincr + 1)

vtk_u = io.VTKFile(domain.comm, "results/results_u.pvd", "w")
vtk_p = io.VTKFile(domain.comm, "results/results_p.pvd", "w")
vtk_sig = io.VTKFile(domain.comm, "results/results_sig.pvd", "w")
for i, t in enumerate(load_steps[1:]):
    f.value[-1] = -1e7 * t
    dt.value = 1/Nincr

    converged, it = problem.solve(newton, print_solution=True)
    ml.x.petsc_vec.copy(ml_n.x.petsc_vec)
    u_out = v.sub(0).collapse()
    u_out.name = "Displacement"
    p_out = v.sub(1).collapse()
    p_out.name = "Pressure"
    stress_out = qmap.project_on("Stress", ("DG",0))
    vtk_u.write_function(u_out, t)
    vtk_p.write_function(p_out, t)
    vtk_sig.write_function(stress_out, t)

vtk_u.close()
vtk_p.close()
vtk_sig.close()