from mpi4py import MPI
import numpy as np
import dolfinx
import scifem

import pytest


@pytest.mark.parametrize(
    "entities_list, set_values",
    [
        ([(1, lambda x: x[0] <= 1.0)], {1}),
        ([(2, lambda x: x[0] <= 1.0)], {2}),
        (
            [
                (1, lambda x: x[0] <= 1.0),
                (2, lambda x: x[0] >= 0.0),
            ],
            {2},
        ),
    ],
)
def test_create_celltags_celltags_all(entities_list, set_values):
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
    cell_tag = scifem.create_entity_markers(mesh, mesh.topology.dim, entities_list=entities_list)

    im = mesh.topology.index_map(mesh.topology.dim)
    assert cell_tag.dim == mesh.topology.dim
    assert cell_tag.indices.shape[0] == im.size_local + im.num_ghosts
    assert cell_tag.values.shape[0] == im.size_local + im.num_ghosts
    assert cell_tag.values.dtype == np.int32
    assert set(cell_tag.values) == set_values


@pytest.mark.parametrize(
    "entities_list, set_values",
    [
        ([(1, lambda x: x[0] <= 1.0, False)], {1}),
        ([(2, lambda x: x[0] <= 1.0, False)], {2}),
        (
            [
                (1, lambda x: x[0] <= 1.0, False),
                (2, lambda x: x[0] >= 0.0, False),
            ],
            {2},
        ),
    ],
)
def test_create_facet_tags_all_on_boundary_False(entities_list, set_values):
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
    facet_tag = scifem.create_entity_markers(
        mesh, mesh.topology.dim - 1, entities_list=entities_list
    )
    im = mesh.topology.index_map(mesh.topology.dim - 1)

    assert facet_tag.dim == mesh.topology.dim - 1
    assert facet_tag.indices.shape[0] == im.size_local + im.num_ghosts
    assert facet_tag.values.shape[0] == im.size_local + im.num_ghosts
    assert facet_tag.values.dtype == np.int32
    assert set(facet_tag.values) == set_values


@pytest.mark.parametrize(
    "entities_list, set_values",
    [
        ([(1, lambda x: x[0] <= 1.0, True)], {1}),
        ([(2, lambda x: x[0] <= 1.0, True)], {2}),
        (
            [
                (1, lambda x: x[0] <= 1.0, True),
                (2, lambda x: x[0] >= 0.0, True),
            ],
            {2},
        ),
    ],
)
def test_create_facet_tags_all_on_boundary_True(entities_list, set_values):
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
    facet_tag = scifem.create_entity_markers(
        mesh, mesh.topology.dim - 1, entities_list=entities_list
    )

    mesh.topology.create_entities(1)
    mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
    facet_indices = dolfinx.mesh.exterior_facet_indices(mesh.topology)
    assert facet_tag.dim == mesh.topology.dim - 1
    assert facet_tag.indices.shape[0] == len(facet_indices)
    assert facet_tag.values.shape[0] == len(facet_indices)
    assert facet_tag.values.dtype == np.int32
    assert set(facet_tag.values) == set_values


@pytest.mark.parametrize(
    "entities_list",
    [
        ([(1, lambda x: x[0] < 0.0)]),
        ([(1, lambda x: x[0] < 0.0)]),
        ([]),
    ],
)
def test_create_celltags_empty(entities_list):
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 3, 3)
    cell_tag = scifem.create_entity_markers(mesh, mesh.topology.dim, entities_list=entities_list)

    assert cell_tag.dim == mesh.topology.dim
    assert cell_tag.indices.shape[0] == 0
    assert cell_tag.values.shape[0] == 0
    assert cell_tag.values.dtype == np.int32
    assert set(cell_tag.values) == set()


@pytest.mark.parametrize("edim", [0, 1, 2, 3])
def test_submesh_meshtags(edim):
    mesh = dolfinx.mesh.create_unit_cube(
        MPI.COMM_WORLD,
        3,
        4,
        7,
        cell_type=dolfinx.cpp.mesh.CellType.tetrahedron,
        ghost_mode=dolfinx.cpp.mesh.GhostMode.shared_facet,
    )

    mesh.topology.create_entities(edim)
    emap = mesh.topology.index_map(edim)

    # Put every second owned cell in submesh
    num_entities_local = emap.size_local
    subset_entities = np.arange(0, num_entities_local, 2, dtype=np.int32)
    # Include ghosts entities
    subset_cells = scifem.reverse_mark_entities(emap, subset_entities)

    submesh, entity_to_parent, vertex_to_parent, _ = dolfinx.mesh.create_submesh(
        mesh, edim, subset_cells
    )

    # Create meshtags on the parent mesh
    for i in range(edim + 1):
        mesh.topology.create_entities(i)
        parent_e_map = mesh.topology.index_map(i)
        num_parent_entities = parent_e_map.size_local + parent_e_map.num_ghosts
        values = parent_e_map.local_range[0] + np.arange(num_parent_entities, dtype=np.int32)
        entity_communicator = dolfinx.la.vector(parent_e_map, 1)
        entity_communicator.array[:] = values
        entity_communicator.scatter_forward()
        parent_tag = dolfinx.mesh.meshtags(
            mesh,
            i,
            np.arange(num_parent_entities, dtype=np.int32),
            entity_communicator.array.astype(np.int32),
        )
        sub_tag, sub_entity_to_parent = scifem.transfer_meshtags_to_submesh(
            parent_tag, submesh, vertex_to_parent, entity_to_parent
        )
        submesh.topology.create_connectivity(i, edim)
        midpoints = dolfinx.mesh.compute_midpoints(submesh, i, sub_tag.indices)
        mesh.topology.create_connectivity(i, mesh.topology.dim)
        parent_midpoints = dolfinx.mesh.compute_midpoints(mesh, i, parent_tag.indices)

        np.testing.assert_allclose(midpoints, parent_midpoints[sub_entity_to_parent])


@pytest.mark.parametrize("codim", [0, 1, 2])
@pytest.mark.parametrize("tdim", [1, 2, 3])
@pytest.mark.parametrize(
    "ghost_mode", [dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet]
)
def test_submesh_creator(codim, tdim, ghost_mode):
    edim = tdim - codim
    if edim < 0:
        pytest.xfail("Codim larger than tdim")

    if tdim == 1:
        mesh = dolfinx.mesh.create_unit_interval(MPI.COMM_WORLD, 27, ghost_mode=ghost_mode)
    elif tdim == 2:
        mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 18, ghost_mode=ghost_mode)
    elif tdim == 3:
        mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 7, 5, 8, ghost_mode=ghost_mode)
    else:
        raise ValueError("Invalid tdim")

    tol = 50 * np.finfo(mesh.geometry.x.dtype).eps

    def first_marker(x):
        return x[0] <= 0.5 + tol

    def second_marker(x):
        return x[tdim - 1] >= 0.6 - tol

    first_val = 2
    second_val = 3
    mesh.topology.create_entities(edim)
    emap = mesh.topology.index_map(edim)

    # Extract submesh no longer accumulates entities, so we mark the by hand
    indicator = dolfinx.la.vector(emap, 1, dtype=np.int32)
    indicator.array[:] = 0
    indicator.array[dolfinx.mesh.locate_entities(mesh, edim, first_marker)] = first_val
    indicator.array[dolfinx.mesh.locate_entities(mesh, edim, second_marker)] = second_val
    indicator.scatter_reverse(dolfinx.la.InsertMode.insert)
    indicator.scatter_forward()

    # Constructor we are testing
    etag = dolfinx.mesh.meshtags(
        mesh, edim, np.arange(emap.size_local + emap.num_ghosts, dtype=np.int32), indicator.array
    )
    submesh, cell_map, vertex_map, node_map, sub_etag = scifem.extract_submesh(
        mesh, etag, (first_val, second_val)
    )

    submap_array = scifem.mesh.get_entity_map(cell_map)
    parent_indices = submap_array[sub_etag.indices]
    np.testing.assert_allclose(sub_etag.values, indicator.array[parent_indices])

    # Create with standard constructor (reference)
    e_comm = dolfinx.la.vector(emap, 1)
    e_comm.array[:] = 0
    e_comm.array[dolfinx.mesh.locate_entities(mesh, edim, first_marker)] = 1
    e_comm.array[dolfinx.mesh.locate_entities(mesh, edim, second_marker)] = 1
    e_comm.scatter_reverse(dolfinx.la.InsertMode.add)
    e_comm.scatter_forward()
    sub_entities = np.flatnonzero(e_comm.array).astype(np.int32)
    ref_submesh, ref_cm, ref_vm, ref_nm = dolfinx.mesh.create_submesh(mesh, edim, sub_entities)

    submap_array = scifem.mesh.get_entity_map(cell_map)
    ref_submap_array = scifem.mesh.get_entity_map(ref_cm)
    np.testing.assert_allclose(submap_array, ref_submap_array)

    subvertexmap_array = scifem.mesh.get_entity_map(vertex_map)
    ref_subvertexmap_array = scifem.mesh.get_entity_map(ref_vm)
    np.testing.assert_allclose(subvertexmap_array, ref_subvertexmap_array)

    np.testing.assert_allclose(node_map, ref_nm)
    assert (
        submesh.topology.index_map(edim).size_local
        == ref_submesh.topology.index_map(edim).size_local
    )
    assert (
        submesh.topology.index_map(edim).size_global
        == ref_submesh.topology.index_map(edim).size_global
    )
    assert (
        submesh.topology.index_map(edim).num_ghosts
        == ref_submesh.topology.index_map(edim).num_ghosts
    )


@pytest.mark.parametrize("tdim", [1, 2, 3])
@pytest.mark.parametrize(
    "ghost_mode", [dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet]
)
@pytest.mark.parametrize("dtype", [np.float64, np.float32])
def test_find_interface(tdim, ghost_mode, dtype):
    comm = MPI.COMM_WORLD
    if tdim == 1:
        mesh = dolfinx.mesh.create_unit_interval(comm, 26, ghost_mode=ghost_mode, dtype=dtype)
    elif tdim == 2:
        mesh = dolfinx.mesh.create_unit_square(comm, 10, 13, ghost_mode=ghost_mode, dtype=dtype)
    elif tdim == 3:
        mesh = dolfinx.mesh.create_unit_cube(comm, 6, 5, 8, ghost_mode=ghost_mode, dtype=dtype)
    else:
        raise ValueError("Invalid tdim")

    tol = 50 * np.finfo(dtype).eps

    def half_entities(x):
        return x[0] <= 0.5 + tol

    cell_map = mesh.topology.index_map(tdim)
    all_cells = np.arange(cell_map.size_local + cell_map.num_ghosts, dtype=np.int32)
    values = np.full_like(all_cells, 1, dtype=np.int32)
    values[dolfinx.mesh.locate_entities(mesh, tdim, half_entities)] = 2
    cell_tags = dolfinx.mesh.meshtags(mesh, tdim, all_cells, values)

    interface = scifem.find_interface(cell_tags, (1,), (2,))

    def ref_interface(x):
        return np.isclose(x[0], 0.5)

    ref_interface_facets = dolfinx.mesh.locate_entities(mesh, tdim - 1, ref_interface)
    mesh.topology.create_connectivity(tdim - 1, tdim)
    f_to_c = mesh.topology.connectivity(tdim - 1, tdim)
    ref_twosided_facets = []
    for f in ref_interface_facets:
        if len(f_to_c.links(f)) == 2:
            ref_twosided_facets.append(f)
    np.testing.assert_allclose(interface, ref_twosided_facets)


@pytest.mark.parametrize("dtype", [np.float64, np.float32])
@pytest.mark.parametrize(
    "ghost_mode", [dolfinx.mesh.GhostMode.none, dolfinx.mesh.GhostMode.shared_facet]
)
@pytest.mark.parametrize(
    "cell_type", [dolfinx.cpp.mesh.CellType.tetrahedron, dolfinx.cpp.mesh.CellType.hexahedron]
)
def test_exterior_boundary_subdomain(dtype, ghost_mode, cell_type):
    comm = MPI.COMM_WORLD
    mesh = dolfinx.mesh.create_unit_cube(
        comm, 13, 7, 11, cell_type=cell_type, ghost_mode=ghost_mode, dtype=dtype
    )

    def center(x):
        return (x[0] - 0.5) ** 2 + (x[1] - 0.5) ** 2 + (x[2] - 0.5) ** 2 <= 0.2**2

    tdim = mesh.topology.dim
    cell_map = mesh.topology.index_map(tdim)
    all_cells = np.arange(cell_map.size_local + cell_map.num_ghosts, dtype=np.int32)
    value_map = dolfinx.la.vector(cell_map, 1, dtype=np.int32)
    value_map.array[:] = 0
    value_map.array[dolfinx.mesh.locate_entities(mesh, tdim, center)] = 1
    value_map.scatter_reverse(dolfinx.la.InsertMode.add)
    value_map.scatter_forward()

    values = np.ones_like(all_cells, dtype=np.int32)
    values[value_map.array > 0] = 2

    cell_tags = dolfinx.mesh.meshtags(mesh, tdim, all_cells, values)

    def boundary_check(mesh, facets, cell_tags, values):
        is_in_subdomain = np.isin(cell_tags.values, np.asarray(values))
        mesh.topology.create_connectivity(mesh.topology.dim - 1, mesh.topology.dim)
        num_facets_local = mesh.topology.index_map(mesh.topology.dim - 1).size_local
        f_to_c = mesh.topology.connectivity(mesh.topology.dim - 1, mesh.topology.dim)
        cell_map = mesh.topology.index_map(mesh.topology.dim)
        owned_facets_in_subdomain = []
        owned_facets_on_outside = []
        ghosted_facets_in_subdomain = []
        ghosted_facets_on_outside = []

        facet_map = mesh.topology.index_map(mesh.topology.dim - 1)

        for facet in facets:
            cells = f_to_c.links(facet)
            if len(cells) != 1:
                assert np.sum(is_in_subdomain[cells]) == 1
            elif len(cells) == 1:
                if facet < num_facets_local:
                    # If connected to only one cell, we send it to the other process
                    # that has it to verify that it is a boundary facet
                    if is_in_subdomain[cells[0]]:
                        owned_facets_in_subdomain.append(facet)
                    else:
                        owned_facets_on_outside.append(facet)
                else:
                    if is_in_subdomain[cells[0]]:
                        ghosted_facets_in_subdomain.append(facet)
                    else:
                        ghosted_facets_on_outside.append(facet)
        # Receive facets from other process that is on subdomain on that process
        facet_vector = dolfinx.la.vector(
            mesh.topology.index_map(mesh.topology.dim - 1), 1, dtype=np.int32
        )
        facet_vector.array[:] = 0
        facet_vector.array[ghosted_facets_in_subdomain] = 1
        facet_vector.scatter_reverse(dolfinx.la.InsertMode.add)

        exterior_facets = dolfinx.mesh.exterior_facet_indices(mesh.topology)
        all_local_exterior_facets = scifem.reverse_mark_entities(facet_map, exterior_facets)

        facets_inside_subdomain = np.flatnonzero(facet_vector.array[:num_facets_local]).astype(
            np.int32
        )
        mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim - 1)
        for facet in facets_inside_subdomain:
            cells = f_to_c.links(facet)
            if facet in all_local_exterior_facets:
                assert is_in_subdomain[cells[0]]
                continue
            if len(cells) == 1:
                assert cells[0] < cell_map.size_local
                assert not is_in_subdomain[cells[0]]
            else:
                assert len(cells) == 2
                assert np.sum(is_in_subdomain[cells]) == 1

        facet_vector.array[:] = 0
        facet_vector.array[ghosted_facets_on_outside] = 1
        facet_vector.scatter_reverse(dolfinx.la.InsertMode.add)
        facets_outside_subdomain = np.flatnonzero(facet_vector.array[:num_facets_local]).astype(
            np.int32
        )
        for facet in facets_outside_subdomain:
            cells = f_to_c.links(facet)
            if facet in all_local_exterior_facets:
                assert is_in_subdomain[cells[0]]
                continue
            if len(cells) == 1:
                assert cells[0] < cell_map.size_local
                assert is_in_subdomain[cells[0]]
            else:
                assert len(cells) == 2
                assert np.sum(is_in_subdomain[cells]) == 1

        facet_vector.array[:] = 0
        facet_vector.array[owned_facets_in_subdomain] = 1
        facet_vector.scatter_forward()
        facets_inside_subdomain = np.flatnonzero(facet_vector.array[num_facets_local:]).astype(
            np.int32
        )
        for facet in facets_inside_subdomain:
            cells = f_to_c.links(num_facets_local + facet)
            if num_facets_local + facet in all_local_exterior_facets:
                assert is_in_subdomain[cells[0]]
                continue

            if len(cells) == 1:
                assert cells[0] <= cell_map.size_local
                assert not is_in_subdomain[cells[0]]
            else:
                assert len(cells) == 2
                assert np.sum(is_in_subdomain[cells]) == 1

        facet_vector.array[:] = 0
        facet_vector.array[owned_facets_on_outside] = 1
        facet_vector.scatter_forward()
        facets_outside_subdomain = np.flatnonzero(facet_vector.array[num_facets_local:]).astype(
            np.int32
        )
        for facet in facets_outside_subdomain:
            cells = f_to_c.links(num_facets_local + facet)
            if num_facets_local + facet in all_local_exterior_facets:
                assert is_in_subdomain[cells[0]]
                continue
            if len(cells) == 1:
                assert cells[0] <= cell_map.size_local
                assert is_in_subdomain[cells[0]]
            else:
                assert len(cells) == 2
                assert np.sum(is_in_subdomain[cells]) == 1

    ext_facets_1 = scifem.compute_subdomain_exterior_facets(mesh, cell_tags, (1,))
    boundary_check(mesh, ext_facets_1, cell_tags, (1,))

    # Exterior facets are only those at the interface
    ext_facets_2 = scifem.compute_subdomain_exterior_facets(mesh, cell_tags, (2,))
    boundary_check(mesh, ext_facets_2, cell_tags, (2,))

    # Exterior facets are only exterior
    ext_facets = scifem.compute_subdomain_exterior_facets(mesh, cell_tags, (1, 2))
    boundary_check(mesh, ext_facets, cell_tags, (1, 2))


@pytest.mark.parametrize(
    "cell_type",
    [
        dolfinx.mesh.CellType.tetrahedron,
        dolfinx.mesh.CellType.hexahedron,
        dolfinx.mesh.CellType.triangle,
        dolfinx.mesh.CellType.quadrilateral,
        dolfinx.mesh.CellType.interval,
    ],
)
@pytest.mark.parametrize("Nx", [6, 8, 11])
def test_compute_interface_data(cell_type: dolfinx.mesh.CellType, Nx: int):
    tdim = dolfinx.mesh.cell_dim(cell_type)
    if tdim == 1:
        mesh = dolfinx.mesh.create_unit_interval(
            MPI.COMM_WORLD, Nx, ghost_mode=dolfinx.mesh.GhostMode.shared_facet
        )
    elif tdim == 2:
        mesh = dolfinx.mesh.create_unit_square(
            MPI.COMM_WORLD,
            Nx,
            8,
            cell_type=cell_type,
            ghost_mode=dolfinx.mesh.GhostMode.shared_facet,
        )
    elif tdim == 3:
        mesh = dolfinx.mesh.create_unit_cube(
            MPI.COMM_WORLD,
            Nx,
            3,
            5,
            cell_type=cell_type,
            ghost_mode=dolfinx.mesh.GhostMode.shared_facet,
        )
    else:
        raise ValueError(f"Invalid {cell_type=}")

    interface_loc = 3 / Nx

    def interface(x):
        return np.isclose(x[0], interface_loc)

    def left_domain(x, tol=100 * np.finfo(mesh.geometry.x.dtype).eps):
        return x[0] <= interface_loc + tol

    val_small = 3
    val_big = 8
    assert val_big > val_small
    tdim = mesh.topology.dim
    cell_map = mesh.topology.index_map(tdim)
    num_cells_local = cell_map.size_local + cell_map.num_ghosts
    values = np.full(num_cells_local, val_small, dtype=np.int32)
    values[dolfinx.mesh.locate_entities(mesh, tdim, left_domain)] = val_big
    cell_tags = dolfinx.mesh.meshtags(
        mesh, tdim, np.arange(num_cells_local, dtype=np.int32), values
    )

    mesh.topology.create_connectivity(tdim - 1, tdim)
    interface_facets = dolfinx.mesh.locate_entities(mesh, tdim - 1, interface)

    interface_data = scifem.compute_interface_data(cell_tags, interface_facets)

    assert np.isin(interface_data[:, 0], cell_tags.find(val_small)).all()
    assert np.isin(interface_data[:, 2], cell_tags.find(val_big)).all()