from mpi4py import MPI
import dolfinx.fem.petsc
import ufl
import numpy as np
import numpy.typing as npt
from enum import Enum
import argparse


class PenaltyMode(Enum):
    BC = 1
    QUADRATIC = 2


def tr_a(v):
    return v("+")


def tr_b(v):
    return v("-")


def left_dofs(x):
    return np.isclose(x[0], 0.0)


def right_dofs(x):
    return np.isclose(x[0], 1.0)


def compute_interface_data(
    cell_tags: dolfinx.mesh.MeshTags, facet_indices: npt.NDArray[np.int32]
) -> npt.NDArray[np.int32]:
    """
    Compute interior facet integrals that are consistently ordered according to the `cell_tags`,
    such that the data `(cell0, facet_idx0, cell1, facet_idx1)` is ordered such that
    `cell_tags[cell0]`<`cell_tags[cell1]`, i.e the cell with the lowest cell marker is considered the
    "+" restriction".

    Args:
        cell_tags: MeshTags that must contain an integer marker for all cells adjacent to the `facet_indices`
        facet_indices: List of facets (local index) that are on the interface.
    Returns:
        The integration data.
    """
    cell_tags.topology.create_connectivity(cell_tags.dim - 1, cell_tags.dim)
    idata = dolfinx.cpp.fem.compute_integration_domains(
        dolfinx.fem.IntegralType.interior_facet,
        cell_tags.topology,
        facet_indices,
    )
    ordered_idata = idata.reshape(-1, 4).copy()
    switch = (
        cell_tags.values[ordered_idata[:, 0]] > cell_tags.values[ordered_idata[:, 2]]
    )
    if True in switch:
        ordered_idata[switch, :] = ordered_idata[switch][:, [2, 3, 0, 1]]
    return ordered_idata


def solve_penalty_problem(
    penalty_mode: PenaltyMode,
    d: float = 0.5,
    Nx: int = 10,
    Ny: int = 10,
    alpha_c: float = 1000.0,
):
    def split(x, tol=1e2 * np.finfo(dolfinx.default_scalar_type).eps):
        """Locator for right subdomain"""
        return x[0] <= d + tol

    mesh = dolfinx.mesh.create_unit_square(
        MPI.COMM_WORLD, Nx, Ny, ghost_mode=dolfinx.mesh.GhostMode.shared_facet
    )

    tdim = mesh.topology.dim
    cell_map = mesh.topology.index_map(tdim)
    num_cells = cell_map.size_local + cell_map.num_ghosts
    values = np.full(num_cells, 2, dtype=np.int32)
    values[dolfinx.mesh.locate_entities(mesh, tdim, split)] = 1
    ct = dolfinx.mesh.meshtags(mesh, tdim, np.arange(num_cells, dtype=np.int32), values)

    omega_A, map_A, v_map_A, n_map_A = dolfinx.mesh.create_submesh(
        mesh, tdim, ct.find(1)
    )
    omega_B, map_B, v_map_B, n_map_B = dolfinx.mesh.create_submesh(
        mesh, tdim, ct.find(2)
    )

    V_A = dolfinx.fem.functionspace(omega_A, ("Lagrange", 1))
    V_B = dolfinx.fem.functionspace(omega_B, ("Lagrange", 1))
    W = ufl.MixedFunctionSpace(*(V_A, V_B))

    u_a = dolfinx.fem.Function(V_A)
    u_b = dolfinx.fem.Function(V_B)
    v_a, v_b = ufl.TestFunctions(W)

    facets_A = dolfinx.mesh.compute_incident_entities(
        mesh.topology, ct.find(1), tdim, tdim - 1
    )
    facets_B = dolfinx.mesh.compute_incident_entities(
        mesh.topology, ct.find(2), tdim, tdim - 1
    )
    gamma = np.intersect1d(facets_A, facets_B)

    dfx = dolfinx.default_scalar_type
    n_sorption = dolfinx.fem.Constant(mesh, dfx(1.0 / 2))
    bc_A = dolfinx.fem.dirichletbc(
        dolfinx.fem.Constant(omega_A, dfx(0.0)),
        dolfinx.fem.locate_dofs_geometrical(V_A, left_dofs),
        V_A,
    )
    bc_B = dolfinx.fem.dirichletbc(
        dolfinx.fem.Constant(omega_B, dfx(n_sorption / d ** (1 / n_sorption - 1))),
        dolfinx.fem.locate_dofs_geometrical(V_B, right_dofs),
        V_B,
    )
    bcs = [bc_A, bc_B]

    idata = compute_interface_data(ct, gamma)

    dGamma = ufl.Measure(
        "dS", domain=mesh, subdomain_data=[(1, idata.flatten())], subdomain_id=1
    )
    dx = ufl.Measure("dx", domain=mesh, subdomain_data=ct)
    dOmegaA = dx(1)
    dOmegaB = dx(2)

    x = ufl.SpatialCoordinate(mesh)
    f_A = -ufl.div(ufl.grad(x[0]))
    f_B = -ufl.div(
        ufl.grad(n_sorption / (d ** (1 / n_sorption - 1)) * x[0] ** (1 / n_sorption))
    )

    F = ufl.inner(ufl.grad(u_a), ufl.grad(v_a)) * dOmegaA
    F += ufl.inner(ufl.grad(u_b), ufl.grad(v_b)) * dOmegaB
    F -= ufl.inner(f_A, v_a) * dOmegaA
    F -= ufl.inner(f_B, v_b) * dOmegaB
    alpha = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(alpha_c))
    K_a = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1.0))
    K_b = dolfinx.fem.Constant(
        mesh,
        dolfinx.default_scalar_type(n_sorption * d ** ((n_sorption - 1) / n_sorption)),
    )

    # Add penalty term for alpha/2 * equation**2 dGamma
    tol = dolfinx.fem.Constant(mesh, 1e2 * np.finfo(dolfinx.default_scalar_type).eps)
    cond_b = ufl.gt(abs(tr_b(u_b)), tol)
    u_b_padded = ufl.conditional(cond_b, tr_b(u_b), tol)
    equation = tr_a(u_a) - K_a * (u_b_padded / K_b) ** n_sorption
    deq_du_a = ufl.diff(equation, u_a) * tr_a(v_a)
    deq_du_b = ufl.diff(equation, u_b) * tr_b(v_b)

    if penalty_mode == PenaltyMode.QUADRATIC:
        F += alpha * equation * (deq_du_a + deq_du_b) * dGamma
    elif penalty_mode == PenaltyMode.BC:
        F += alpha * equation * (tr_a(v_a) - tr_b(v_b)) * dGamma
    else:
        raise ValueError(f"Unknown penalty mode: {penalty_mode}")
    F_blocked = ufl.extract_blocks(F)

    petsc_options = {
        "snes_type": "newtonls",
        "snes_linesearch_type": "none",
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
        "snes_monitor": None,
        "snes_error_if_not_converged": True,
        "ksp_error_if_not_converged": True,
        "ksp_monitor": None,
        "snes_atol": 1e-10,
        "snes_rtol": 1e-10,
    }

    solver = dolfinx.fem.petsc.NonlinearProblem(
        F_blocked,
        [u_a, u_b],
        bcs=bcs,
        petsc_options=petsc_options,
        petsc_options_prefix="penalty_solver_",
        entity_maps=[map_A, map_B],
        kind="mpi",
    )
    solver.solve()

    with dolfinx.io.VTXWriter(omega_A.comm, "u_A.bp", [u_a]) as bp:
        bp.write(0.0)

    with dolfinx.io.VTXWriter(omega_B.comm, "u_B.bp", [u_b]) as bp:
        bp.write(0.0)

    ub_ex = n_sorption * x[0] ** (1 / n_sorption) / d ** (1 / n_sorption - 1)
    ua_ex = x[0]

    u_b_error = ufl.inner(u_b - ub_ex, u_b - ub_ex) * dOmegaB
    L2_b = dolfinx.fem.form(u_b_error, entity_maps=[map_B])
    print("L2 error u_b:", np.sqrt(dolfinx.fem.assemble_scalar(L2_b)))
    u_a_error = ufl.inner(u_a - ua_ex, u_a - ua_ex) * dOmegaA
    L2_a = dolfinx.fem.form(u_a_error, entity_maps=[map_A])
    print("L2 error u_a:", np.sqrt(dolfinx.fem.assemble_scalar(L2_a)))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Solve penalty problem.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--penalty_mode",
        type=str,
        choices=[e.name for e in PenaltyMode],
        default=PenaltyMode.QUADRATIC.name,
        help="Penalty mode to use: BC or QUADRATIC.",
    )
    parser.add_argument(
        "--Nx", type=int, default=10, help="Number of elements in the x direction."
    )
    parser.add_argument(
        "--Ny", type=int, default=10, help="Number of elements in the y direction."
    )
    parser.add_argument(
        "--d",
        type=float,
        default=0.1,
        help="Parameter d for the penalty problem.",
    )
    parser.add_argument(
        "--alpha", type=float, default=1000.0, help="Penalty parameter alpha."
    )
    args = parser.parse_args()

    penalty_mode = PenaltyMode[args.penalty_mode]
    solve_penalty_problem(
        penalty_mode=penalty_mode, d=args.d, Nx=args.Nx, Ny=args.Ny, alpha_c=args.alpha
    )