import os
import pytest
from pyvista.core.utilities import from_meshio
from pyhip.commands.mesh_idcard import extract_hdf_meshinfo
import yamio


HIP_NAME = {"tetra": "Tetrahedra", "hexahedron": "Hexahedra"}

ATOL = 1e-8


NOT_SUPPORTED = {
    "xml": set(["quad", "hexahedron"]),
}


@pytest.mark.parametrize("elem_type", ["tetra", "hexahedron"])
def test_thickshell_hip_export(tmpdir, param_thickaxishell, elem_type):
    """
    Notes:
        If mesh is read by hip, it is assumed to be working properly.

        Note verification of negative volumes is already done in HipWriter.
    """
    expected_patch_labels = ["Bottom", "Top", "Left", "Right", "Back", "Front"]

    filename = "shell.mesh.h5"
    _, shell = param_thickaxishell

    mesh_filename = os.path.join(tmpdir, filename)
    shell.export_mesh(mesh_filename, elem_type)

    mesh_info, _, _, patch_labels = extract_hdf_meshinfo(mesh_filename)

    # check element type
    assert list(mesh_info["Number of cells"].keys())[0] == HIP_NAME[elem_type]
    assert mesh_info["Number of boundary nodes"] > 0
    for patch_label, expected_patch_label in zip(patch_labels, expected_patch_labels):
        assert patch_label == expected_patch_label


@pytest.mark.parametrize("fmt", ["xdmf", "xml"])
@pytest.mark.parametrize("elem_type", ["tetra", "hexahedron"])
def test_thickshell_export(tmpdir, param_thickaxishell, elem_type, fmt):
    """
    Notes:
        Patches are lost in these formats.

        Writes and reads mesh. Uses pyvista to check volumes.
    """
    if elem_type in NOT_SUPPORTED.get(fmt, []):
        return

    filename = f"shell.{fmt}"
    _, shell = param_thickaxishell
    init_mesh = shell.get_mesh(elem_type)

    # write in given fmt
    mesh_filename = os.path.join(tmpdir, filename)
    shell.export_mesh(mesh_filename, elem_type)

    # read mesh
    mesh = yamio.read(mesh_filename)

    # compare meshes
    # TODO: update when yamio has new functionaties
    mesh = yamio.Mesh(mesh.points, mesh.cells)
    init_mesh.bnd_patches = {}
    assert mesh == init_mesh

    # check areas (if positive)
    cell_volumes = from_meshio(mesh).compute_cell_sizes().cell_data["Volume"]
    assert all(cell_volumes > 0.0)


@pytest.mark.parametrize("fmt", ["xdmf", "xml"])
@pytest.mark.parametrize("elem_type", ["triangle", "quad"])
def test_shell2d_export(tmpdir, shell_2d, elem_type, fmt):
    """
    Notes:
        Writes and reads mesh, looking to the points and cells.
    """
    if elem_type in NOT_SUPPORTED.get(fmt, []):
        return

    shell_type, shell = shell_2d
    if shell_type in ["torchshell"]:  # torch has null areas
        return

    filename = f"shell.{fmt}"

    # write in given fmt
    mesh_filename = os.path.join(tmpdir, filename)
    init_mesh = shell.get_mesh(elem_type)
    yamio.write(mesh_filename, init_mesh)

    # read back
    mesh = yamio.read(mesh_filename)

    # compare meshes
    # TODO: update when yamio has new functionaties
    mesh = yamio.Mesh(mesh.points, mesh.cells)
    init_mesh.bnd_patches = {}
    assert mesh == init_mesh

    # check areas (if positive)
    cell_areas = from_meshio(mesh).compute_cell_sizes().cell_data["Area"]
    assert all(cell_areas > 0.0)
