import importlib

from mpi4py import MPI

import dolfinx
import numpy as np
import pytest
import ufl

from io4dolfinx import FileMode, read_mesh, write_mesh


def get_hdf5_version():
    """Get the HDF5 library version found on the system.

    Note: This function first tries to get the version via h5py.
    If h5py is not installed, it attempts to load the HDF5 shared library
    directly using ctypes to query the version.
    """
    import ctypes
    from ctypes import byref, c_uint

    from packaging.version import parse as _v

    try:
        import h5py

        return _v(h5py.version.hdf5_version)
    except ImportError:
        try:
            # Try to load default HDF5 library names
            # If adios2 is already imported, the symbols might be globally available
            try:
                libhdf5 = ctypes.CDLL("libhdf5.so")
            except OSError:
                try:
                    libhdf5 = ctypes.CDLL("libhdf5.dylib")  # MacOS
                except OSError:
                    print("\n[Method 2] Could not load libhdf5 shared library directly.")
                    libhdf5 = None

            if libhdf5:
                major = c_uint()
                minor = c_uint()
                rel = c_uint()

                # Call the C function
                libhdf5.H5get_libversion(byref(major), byref(minor), byref(rel))
                return _v(f"{major.value}.{minor.value}.{rel.value}")
        except Exception:
            pass

    except Exception:
        pass

    raise RuntimeError("Failed to get HDF5 version")


@pytest.mark.parametrize(
    "backend, encoder, suffix",
    [
        ("adios2", "BP4", ".bp"),
        ("adios2", "HDF5", ".h5"),
        ("adios2", "BP5", ".bp"),
        ("h5py", "HDF5", ".h5"),
    ],
)
@pytest.mark.parametrize(
    "ghost_mode", [dolfinx.mesh.GhostMode.shared_facet, dolfinx.mesh.GhostMode.none]
)
@pytest.mark.parametrize("store_partition", [True, False])
def test_mesh_read_writer(backend, encoder, suffix, ghost_mode, tmp_path, store_partition):
    N = 7

    if backend == "adios2" and encoder == "HDF5" and get_hdf5_version().major >= 2:
        pytest.skip("HDF5 version >= 2 is not supported due to ADIOS2 limitations.")
    try:
        importlib.import_module(backend)
    except ModuleNotFoundError:
        pytest.skip(f"{backend} not installed")

    # Consistent tmp dir across processes
    fname = MPI.COMM_WORLD.bcast(tmp_path, root=0)
    file = fname / f"{backend}_mesh_{encoder}_{store_partition}"
    xdmf_file = fname / f"xdmf_mesh_{encoder}_{ghost_mode}_{store_partition}"

    mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N, ghost_mode=ghost_mode)

    backend_args = None
    if backend == "adios2":
        backend_args = {"engine": encoder}
    write_mesh(
        file.with_suffix(suffix),
        mesh,
        store_partition_info=store_partition,
        backend_args=backend_args,
        backend=backend,
    )

    mesh.comm.Barrier()
    with dolfinx.io.XDMFFile(mesh.comm, xdmf_file.with_suffix(".xdmf"), "w") as xdmf:
        xdmf.write_mesh(mesh)
    mesh.comm.Barrier()

    in_mesh = read_mesh(
        file.with_suffix(suffix),
        MPI.COMM_WORLD,
        ghost_mode=ghost_mode,
        read_from_partition=store_partition,
        backend_args=backend_args,
        backend=backend,
    )
    in_mesh.comm.Barrier()
    if store_partition:

        def compute_distance_matrix(points_A, points_B, tol=1e-12):
            points_A_e = np.expand_dims(points_A, 1)
            points_B_e = np.expand_dims(points_B, 0)
            distances = np.sum(np.square(points_A_e - points_B_e), axis=2)
            return distances < tol

        cell_map = mesh.topology.index_map(mesh.topology.dim)
        new_cell_map = in_mesh.topology.index_map(in_mesh.topology.dim)
        assert cell_map.size_local == new_cell_map.size_local
        assert cell_map.num_ghosts == new_cell_map.num_ghosts
        mesh.topology.create_connectivity(mesh.topology.dim, mesh.topology.dim)
        midpoints = dolfinx.mesh.compute_midpoints(
            mesh,
            mesh.topology.dim,
            np.arange(cell_map.size_local + cell_map.num_ghosts, dtype=np.int32),
        )
        in_mesh.topology.create_connectivity(in_mesh.topology.dim, in_mesh.topology.dim)
        new_midpoints = dolfinx.mesh.compute_midpoints(
            in_mesh,
            in_mesh.topology.dim,
            np.arange(new_cell_map.size_local + new_cell_map.num_ghosts, dtype=np.int32),
        )
        # Check that all points in owned by initial mesh is owned by the new mesh
        # (might be locally reordered)
        owned_distances = compute_distance_matrix(
            midpoints[: cell_map.size_local], new_midpoints[: new_cell_map.size_local]
        )
        np.testing.assert_allclose(np.sum(owned_distances, axis=1), 1)
        # Check that all points that are ghosted in original mesh is ghosted on the
        # same process in the new mesh
        ghost_distances = compute_distance_matrix(
            midpoints[cell_map.size_local :], new_midpoints[new_cell_map.size_local :]
        )
        np.testing.assert_allclose(np.sum(ghost_distances, axis=1), 1)

    mesh.comm.Barrier()

    with dolfinx.io.XDMFFile(mesh.comm, xdmf_file.with_suffix(".xdmf"), "r") as xdmf:
        mesh_xdmf = xdmf.read_mesh(ghost_mode=ghost_mode)

    for i in range(mesh.topology.dim + 1):
        mesh.topology.create_entities(i)
        mesh_xdmf.topology.create_entities(i)
        in_mesh.topology.create_entities(i)
        assert (
            mesh_xdmf.topology.index_map(i).size_global == in_mesh.topology.index_map(i).size_global
        )

    # Check that integration over different entities are consistent
    measures = (
        [ufl.ds, ufl.dx] if ghost_mode is dolfinx.mesh.GhostMode.none else [ufl.ds, ufl.dS, ufl.dx]
    )
    for measure in measures:
        form = dolfinx.fem.form(1 * measure(domain=in_mesh))
        c_adios = dolfinx.fem.assemble_scalar(form)
        c_ref = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh)))
        c_xdmf = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh_xdmf)))
        assert np.isclose(
            in_mesh.comm.allreduce(c_adios, MPI.SUM),
            mesh.comm.allreduce(c_xdmf, MPI.SUM),
        )
        assert np.isclose(
            in_mesh.comm.allreduce(c_adios, MPI.SUM),
            mesh.comm.allreduce(c_ref, MPI.SUM),
        )


@pytest.mark.parametrize(
    "backend, encoder, suffix",
    [("adios2", "BP4", ".bp"), ("adios2", "BP5", ".bp"), ("h5py", None, ".h5")],
)
@pytest.mark.parametrize(
    "ghost_mode", [dolfinx.mesh.GhostMode.shared_facet, dolfinx.mesh.GhostMode.none]
)
@pytest.mark.parametrize("store_partition", [True, False])
def test_timedep_mesh(encoder, backend, suffix, ghost_mode, tmp_path, store_partition):
    try:
        importlib.import_module(backend)
    except ModuleNotFoundError:
        pytest.skip(f"{backend} not installed")

    # Currently unsupported, unclear why ("HDF5", ".h5"),
    N = 13
    # Consistent tmp dir across processes
    fname = MPI.COMM_WORLD.bcast(tmp_path, root=0)
    file = fname / f"adios_time_dep_mesh_{encoder}"
    mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, N, N, N, ghost_mode=ghost_mode)

    def u(x):
        return np.asarray([x[0] + 0.1 * np.sin(x[1]), 0.2 * np.cos(x[1]), x[2]])

    backend_args = None
    if backend == "adios2":
        backend_args = {"engine": encoder}
    write_mesh(
        file.with_suffix(suffix),
        mesh,
        mode=FileMode.write,
        time=0.0,
        store_partition_info=store_partition,
        backend_args=backend_args,
        backend=backend,
    )
    delta_x = u(mesh.geometry.x.T).T
    mesh.geometry.x[:] += delta_x
    write_mesh(
        file.with_suffix(suffix),
        mesh,
        mode=FileMode.append,
        time=3.0,
        backend_args=backend_args,
        backend=backend,
    )
    mesh.geometry.x[:] -= delta_x

    mesh_first = read_mesh(
        file.with_suffix(suffix),
        MPI.COMM_WORLD,
        ghost_mode,
        time=0.0,
        read_from_partition=store_partition,
        backend_args=backend_args,
        backend=backend,
    )
    mesh_first.comm.Barrier()

    # Check that integration over different entities are consistent
    measures = [ufl.ds, ufl.dx]
    if ghost_mode == dolfinx.mesh.GhostMode.shared_facet:
        measures.append(ufl.dx)
    for measure in measures:
        # try:
        c_adios = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh_first)))
        c_ref = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh)))
        assert np.isclose(
            mesh_first.comm.allreduce(c_adios, MPI.SUM),
            mesh.comm.allreduce(c_ref, MPI.SUM),
        )

    mesh.geometry.x[:] += delta_x
    mesh_second = read_mesh(
        file.with_suffix(suffix),
        MPI.COMM_WORLD,
        ghost_mode,
        time=3.0,
        read_from_partition=store_partition,
        backend_args=backend_args,
        backend=backend,
    )
    mesh_second.comm.Barrier()
    measures = [ufl.ds, ufl.dx]
    if ghost_mode == dolfinx.mesh.GhostMode.shared_facet:
        measures.append(ufl.dx)
    for measure in measures:
        c_adios = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh_second)))
        c_ref = dolfinx.fem.assemble_scalar(dolfinx.fem.form(1 * measure(domain=mesh)))
        assert np.isclose(
            mesh_second.comm.allreduce(c_adios, MPI.SUM),
            mesh.comm.allreduce(c_ref, MPI.SUM),
        )
