if args.radius > 0.8 / args.time_steps:
                args.radius = 0.8 / args.time_steps

import argparse
import sys

import numpy as np
import ufl
from dolfinx import default_scalar_type, log
from dolfinx.common import timed, Timer, TimingType, list_timings, timing
from dolfinx.fem import (Constant, dirichletbc, form, Function, Expression, FunctionSpace,
                         VectorFunctionSpace, locate_dofs_topological)
from dolfinx.fem.petsc import create_vector, assemble_vector, assemble_matrix, set_bc, apply_lifting
from dolfinx.graph import adjacencylist
from dolfinx.io import XDMFFile, VTXWriter
from dolfinx.mesh import locate_entities_boundary, GhostMode, meshtags
from mpi4py import MPI
from petsc4py.PETSc import InsertMode, ScatterMode  # type: ignore

from dolfinx_contact.helpers import (epsilon, lame_parameters, sigma_func,
                                     weak_dirichlet)
from dolfinx_contact.meshing import (convert_mesh, create_box_mesh_2D,
                                     create_box_mesh_3D,
                                     create_circle_circle_mesh,
                                     create_circle_plane_mesh,
                                     create_cylinder_cylinder_mesh,
                                     create_sphere_plane_mesh)
from dolfinx_contact.general_contact.contact_problem import ContactProblem, FrictionLaw
from dolfinx_contact.parallel_mesh_ghosting import create_contact_mesh
from dolfinx_contact.helpers import rigid_motions_nullspace_subdomains
from dolfinx_contact.newton_solver import NewtonSolver
from dolfinx_contact.cpp import ContactMode
import os
from typing import Union
from dolfinx import cpp

if __name__ == "__main__":
    desc = "Nitsche's method for two elastic bodies using custom assemblers"
    parser = argparse.ArgumentParser(description=desc,
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("--theta", default=1., type=float, dest="theta",
                        help="Theta parameter for Nitsche, 1 symmetric, -1 skew symmetric, 0 Penalty-like",
                        choices=[1., -1., 0.])
    parser.add_argument("--gamma", default=10, type=float, dest="gamma",
                        help="Coercivity/Stabilization parameter for Nitsche condition")
    parser.add_argument("--quadrature", default=5, type=int, dest="q_degree",
                        help="Quadrature degree used for contact integrals")
    parser.add_argument("--problem", default=1, type=int, dest="problem",
                        help="Which problem to solve: 1. Flat surfaces, 2. One curved surface, 3. Two curved surfaces",
                        choices=[1, 2, 3])
    parser.add_argument("--order", default=1, type=int, dest="order",
                        help="Order of mesh geometry", choices=[1, 2, 3])
    _3D = parser.add_mutually_exclusive_group(required=False)
    _3D.add_argument('--3D', dest='threed', action='store_true',
                     help="Use 3D mesh", default=False)
    
    
    _timing = parser.add_mutually_exclusive_group(required=False)
    _timing.add_argument('--timing', dest='timing', action='store_true',
                         help="List timings", default=False)
    _ksp = parser.add_mutually_exclusive_group(required=False)
    _ksp.add_argument('--ksp-view', dest='ksp', action='store_true',
                      help="List ksp options", default=False)
    _simplex = parser.add_mutually_exclusive_group(required=False)
    _simplex.add_argument('--simplex', dest='simplex', action='store_true',
                          help="Use triangle/tet mesh", default=False)
    _strain = parser.add_mutually_exclusive_group(required=False)
    _strain.add_argument('--strain', dest='plane_strain', action='store_true',
                         help="Use plane strain formulation", default=False)
    parser.add_argument("--E", default=1e3, type=np.float64, dest="E",
                        help="Youngs modulus of material")
    parser.add_argument("--nu", default=0.1, type=np.float64, dest="nu", help="Poisson's ratio")
    parser.add_argument("--disp", default=0.2, type=np.float64, dest="disp",
                        help="Displacement BC in negative y direction")
    parser.add_argument("--radius", default=0.5, type=np.float64, dest="radius",
                        help="Search radius for ray-tracing")
    parser.add_argument("--load_steps", default=4, type=np.int32, dest="nload_steps",
                        help="Number of steps for gradual loading")
    parser.add_argument("--res", default=0.1, type=np.float64, dest="res",
                        help="Mesh resolution")
    parser.add_argument("--friction", default=0.0, type=np.float64, dest="fric",
                        help="Friction coefficient for Tresca friction")
    parser.add_argument("--outfile", type=str, default=None, required=False,
                        help="File for appending results", dest="outfile")
    _lifting = parser.add_mutually_exclusive_group(required=False)
    _lifting.add_argument('--lifting', dest='lifting', action='store_true',
                          help="Apply lifting (strong enforcement of Dirichlet condition",
                          default=False)
    _raytracing = parser.add_mutually_exclusive_group(required=False)
    _raytracing.add_argument('--raytracing', dest='raytracing', action='store_true',
                             help="Use raytracing for contact search.",
                             default=False)
    _coulomb = parser.add_mutually_exclusive_group(required=False)
    _coulomb.add_argument('--coulomb', dest='coulomb', action='store_true',
                          help="Use coulomb friction kernel. This requires --friction=[nonzero float value].",
                          default=False)

    # Parse input arguments or set to defualt values
    args = parser.parse_args()
    # Current formulation uses bilateral contact
    threed = args.threed
    problem = args.problem
    nload_steps = args.nload_steps
    simplex = args.simplex
    mesh_dir = "code_utils/tutoriales/results/meshes"
    if not os.path.exists(mesh_dir):
        os.makedirs(mesh_dir)
    triangle_ext = {1: "", 2: "6", 3: "10"}
    tetra_ext = {1: "", 2: "10", 3: "20"}
    hex_ext = {1: "", 2: "27"}
    quad_ext = {1: "", 2: "9", 3: "16"}
    line_ext = {1: "", 2: "3", 3: "4"}
    if args.order > 2:
        raise NotImplementedError("More work in DOLFINx (SubMesh) required for this to work.")
    # Load mesh and create identifier functions for the top (Displacement condition)
    # and the bottom (contact condition)

    if threed:
        displacement = np.array([[0, 0, -args.disp], [0, 0, 0]])
        if problem == 1:
            outname = "code_utils/tutoriales/results/tu8/problem2_3D_simplex" if simplex else "code_utils/tutoriales/results/tu8/problem2_3D_hex"
            fname = f"{mesh_dir}/sphere"
            create_sphere_plane_mesh(filename=f"{fname}.msh", order=args.order, res=args.res)
            convert_mesh(fname, fname, gdim=3)
            with XDMFFile(MPI.COMM_WORLD, f"{fname}.xdmf", "r") as xdmf:
                mesh = xdmf.read_mesh()
                domain_marker = xdmf.read_meshtags(mesh, name="cell_marker")
                tdim = mesh.topology.dim
                mesh.topology.create_connectivity(tdim - 1, tdim)
                facet_marker = xdmf.read_meshtags(mesh, name="facet_marker")
            dirichlet_bdy_1 = 2
            contact_bdy_1 = 1
            contact_bdy_2 = 8
            dirichlet_bdy_2 = 7

    else:
        displacement = np.array([[0, -args.disp], [0, 0]])
        if problem == 1:
            outname = "code_utils/tutoriales/results/tu8/problem2_2D_simplex" if simplex else "code_utils/tutoriales/results/tu8/problem2_2D_quads"
            fname = f"{mesh_dir}/problem2_2D_simplex" if simplex else f"{mesh_dir}/problem2_2D_quads"
            create_circle_plane_mesh(filename=f"{fname}.msh", quads=not simplex,
                                     res=args.res, order=args.order, r=0.3, gap=0.1, height=0.1, length=1.0)
            convert_mesh(fname, f"{fname}.xdmf", gdim=2)

            with XDMFFile(MPI.COMM_WORLD, f"{fname}.xdmf", "r") as xdmf:
                mesh = xdmf.read_mesh()
                domain_marker = xdmf.read_meshtags(mesh, name="cell_marker")
                tdim = mesh.topology.dim
                mesh.topology.create_connectivity(tdim - 1, tdim)
                facet_marker = xdmf.read_meshtags(mesh, name="facet_marker")
            dirichlet_bdy_1 = 8
            contact_bdy_1 = 10
            contact_bdy_2 = 6
            dirichlet_bdy_2 = 4

    if mesh.comm.size > 1:
        mesh, facet_marker, domain_marker = create_contact_mesh(
            mesh, facet_marker, domain_marker, [contact_bdy_1, contact_bdy_2], 2.0)

    ncells = mesh.topology.index_map(tdim).size_local
    indices = np.array(range(ncells), dtype=np.int32)
    values = mesh.comm.rank * np.ones(ncells, dtype=np.int32)
    process_marker = meshtags(mesh, tdim, indices, values)
    process_marker.name = "process_marker"
    domain_marker.name = "cell_marker"
    facet_marker.name = "facet_marker"
    with XDMFFile(mesh.comm, f"{mesh_dir}/test.xdmf", "w") as xdmf:
        xdmf.write_mesh(mesh)
        xdmf.write_meshtags(domain_marker, mesh.geometry)
        xdmf.write_meshtags(facet_marker, mesh.geometry)

    # Solver options
    ksp_tol = 1e-10
    newton_tol = 1e-7
    newton_options = {"relaxation_parameter": 1,
                      "atol": newton_tol,
                      "rtol": newton_tol,
                      "convergence_criterion": "residual",
                      "max_it": 200,
                      "error_on_nonconvergence": True}
    # petsc_options = {"ksp_type": "preonly", "pc_type": "lu"}
    petsc_options = {
        "matptap_via": "scalable",
        "ksp_type": "cg",
        "ksp_rtol": ksp_tol,
        "ksp_atol": ksp_tol,
        "pc_type": "gamg",
        "pc_mg_levels": 3,
        "pc_mg_cycles": 1,   # 1 is v, 2 is w
        "mg_levels_ksp_type": "chebyshev",
        "mg_levels_pc_type": "jacobi",
        "pc_gamg_type": "agg",
        "pc_gamg_coarse_eq_limit": 1000,
        "pc_gamg_agg_nsmooths": 1,
        "pc_gamg_threshold": 0.015,
        "pc_gamg_square_graph": 2,
        "pc_gamg_reuse_interpolation": True,
        "ksp_norm_type": "unpreconditioned"
    }
    # Pack mesh data for Nitsche solver
    dirichlet_vals = [dirichlet_bdy_1, dirichlet_bdy_2]
    contact = [(1, 0), (0, 1)]
    data = np.array([contact_bdy_1, contact_bdy_2], dtype=np.int32)
    offsets = np.array([0, 2], dtype=np.int32)
    surfaces = adjacencylist(data, offsets)

    # Function, TestFunction, TrialFunction and measures
    V = ufl.VectorElement("P", mesh.ufl_cell(), 1)
    Q = ufl.TensorElement("DG", mesh.ufl_cell(), 1)    # , shape = (3,3)
    W = FunctionSpace(mesh, ufl.MixedElement([V, Q]))

    s = Function(W)
    (u, epsv) = ufl.split(s)
    s_old = Function(W)
    (du, epsv_old) = ufl.split(s_old)
    s_ = ufl.TestFunction(W)
    (v, epsv_) = ufl.split(s_)
    s_t = ufl.TrialFunction(W)  
    (w, epsv_t) = ufl.split(s_t)

    # V = VectorFunctionSpace(mesh, ("Lagrange", args.order))
    # u = Function(V)
    # du = Function(V)
    # v = ufl.TestFunction(V)
    # w = ufl.TrialFunction(V)
    dx = ufl.Measure("dx", domain=mesh, subdomain_data=domain_marker)
    ds = ufl.Measure("ds", domain=mesh, subdomain_data=facet_marker)

    # Compute lame parameters
    E = args.E
    nu = args.nu
    mu_func, lambda_func = lame_parameters(args.plane_strain)
    mu = mu_func(E, nu)
    lmbda = lambda_func(E, nu)
    sigma = sigma_func(mu, lmbda)

    # Create variational form without contact contributions
    F = ufl.inner(sigma(u), epsilon(v)) * dx

    # Nitsche parameters
    gamma = args.gamma
    theta = args.theta
    gdim = mesh.geometry.dim
    bcs = []
    bc_fns = []
    for k in range(displacement.shape[0]):
        d = displacement[k, :]
        tag = dirichlet_vals[k]
        g = Constant(mesh, default_scalar_type(tuple(d[i] for i in range(gdim))))
        bc_fns.append(g)

        def weak_dirichlet(F: ufl.Form, u: Function,
                        f: Union[Function, Constant], sigma, gamma, theta, ds):
            V = s.function_space
            # V = W.sub(0)  # I also use this
            # v = F.arguments()[0]

            mesh = V.mesh
            V2 = FunctionSpace(mesh, ("DG", 0))
            tdim = mesh.topology.dim
            ncells = mesh.topology.index_map(tdim).size_local
            h = Function(V2)
            h_vals = cpp.mesh.h(mesh._cpp_object, mesh.topology.dim, np.arange(0, ncells, dtype=np.int32))
            h.x.array[:ncells] = h_vals[:]
            n = ufl.FacetNormal(mesh)
            
            F += - ufl.inner(sigma(u) * n, v) * ds\
                - theta * ufl.inner(sigma(v) * n, u - f) * \
                ds + gamma / h * ufl.inner(u - f, v) * ds
            return F

        F = weak_dirichlet(F, u, g, sigma, E * gamma * args.order**2, theta, ds(tag))

    F = ufl.replace(F, {u: u + du})
    J = ufl.derivative(F, du, w)

    # compiler options to improve performance
    cffi_options = ["-Ofast", "-march=native"]
    jit_options = {"cffi_extra_compile_args": cffi_options,
                   "cffi_libraries": ["m"]}
    # compiled forms for rhs and tangen system
    F_compiled = form(F, jit_options=jit_options)
    J_compiled = form(J, jit_options=jit_options)

    log.set_log_level(log.LogLevel.WARNING)
    num_newton_its = np.zeros(nload_steps, dtype=int)
    num_krylov_its = np.zeros(nload_steps, dtype=int)
    newton_time = np.zeros(nload_steps, dtype=np.float64)

    solver_outfile = args.outfile if args.ksp else None

    V0 = FunctionSpace(mesh, ("DG", 0))
    mu0 = Function(V0)
    lmbda0 = Function(V0)
    fric = Function(V0)
    mu0.interpolate(lambda x: np.full((1, x.shape[1]), mu))
    lmbda0.interpolate(lambda x: np.full((1, x.shape[1]), lmbda))
    fric.interpolate(lambda x: np.full((1, x.shape[1]), args.fric))

    if args.raytracing:
        search_mode = [ContactMode.Raytracing for _ in range(len(contact))]
    else:
        search_mode = [ContactMode.ClosestPoint for _ in range(len(contact))]

    # create contact solver
    contact_problem = ContactProblem([facet_marker], surfaces, contact, mesh, args.q_degree, search_mode)
    if args.coulomb:
        friction_law = FrictionLaw.Coulomb
    else:
        friction_law = FrictionLaw.Tresca
    contact_problem.generate_contact_data(friction_law, W.sub(0), {"u": u, "du": du, "mu": mu0,
                                                            "lambda": lmbda0, "fric": fric},
                                          E * gamma * args.order**2, theta)

    # define functions for newton solver
    def compute_coefficients(x, coeffs):
        size_local = V.dofmap.index_map.size_local
        bs = V.dofmap.index_map_bs
        du.x.array[:size_local * bs] = x.array_r[:size_local * bs]
        du.x.scatter_forward()
        contact_problem.update_contact_data(du)

    @timed("~Contact: Assemble residual")
    def compute_residual(x, b, coeffs):
        b.zeroEntries()
        b.ghostUpdate(addv=InsertMode.INSERT,
                      mode=ScatterMode.FORWARD)
        with Timer("~~Contact: Contact contributions (in assemble vector)"):
            contact_problem.assemble_vector(b, V)
        with Timer("~~Contact: Standard contributions (in assemble vector)"):
            assemble_vector(b, F_compiled)

        # Apply boundary condition
        if len(bcs) > 0:
            apply_lifting(
                b, [J_compiled], bcs=[bcs], x0=[x], scale=-1.0)
        b.ghostUpdate(addv=InsertMode.ADD,
                      mode=ScatterMode.REVERSE)
        if len(bcs) > 0:
            set_bc(b, bcs, x, -1.0)

    @timed("~Contact: Assemble matrix")
    def compute_jacobian_matrix(x, a_mat, coeffs):
        a_mat.zeroEntries()
        with Timer("~~Contact: Contact contributions (in assemble matrix)"):
            contact_problem.assemble_matrix(a_mat, V)
        with Timer("~~Contact: Standard contributions (in assemble matrix)"):
            assemble_matrix(a_mat, J_compiled, bcs=bcs)
        a_mat.assemble()

    # create vector and matrix
    a_mat = contact_problem.create_matrix(J_compiled)
    b = create_vector(F_compiled)

    # Set up snes solver for nonlinear solver
    newton_solver = NewtonSolver(mesh.comm, a_mat, b, contact_problem.coeffs)
    # Set matrix-vector computations
    newton_solver.set_residual(compute_residual)
    newton_solver.set_jacobian(compute_jacobian_matrix)
    newton_solver.set_coefficients(compute_coefficients)

    # Set rigid motion nullspace
    null_space = rigid_motions_nullspace_subdomains(V, domain_marker, np.unique(
        domain_marker.values), num_domains=len(np.unique(domain_marker.values)))
    newton_solver.A.setNearNullSpace(null_space)

    # Set Newton solver options
    newton_solver.set_newton_options(newton_options)

    # Set Krylov solver options
    newton_solver.set_krylov_options(petsc_options)

    # initialise vtx writer
    u.name = "u"
    vtx = VTXWriter(mesh.comm, f"{outname}_{nload_steps}_step.bp", [u], "bp4")
    vtx.write(0)
    for i in range(nload_steps):
        for k, g in enumerate(bc_fns):
            if len(bcs) > 0:
                g.value[:] = displacement[k, :] / nload_steps
            else:
                g.value[:] = (i + 1) * displacement[k, :] / nload_steps
        timing_str = f"~Contact: {i+1} Newton Solver"
        with Timer(timing_str):
            n, converged = newton_solver.solve(du, write_solution=True)
        num_newton_its[i] = n

        du.x.scatter_forward()
        u.x.array[:] += du.x.array[:]
        contact_problem.update_contact_detection(u)
        a_mat = contact_problem.create_matrix(J_compiled)
        a_mat.setNearNullSpace(null_space)
        newton_solver.set_petsc_matrix(a_mat)
        du.x.array[:] = 0.1 * du.x.array[:]
        contact_problem.update_contact_data(du)
        vtx.write(i + 1)
    vtx.close()