"""Tests for the mesh submodule: definition and transformation of base meshes."""

import pytest

import numpy as np
from numpy.typing import NDArray

import capytaine as cpt
from capytaine.meshes.meshes import Mesh as OldMesh
from capytaine.meshes.properties import clustering
from capytaine.meshes.predefined import mesh_sphere, mesh_horizontal_cylinder
from capytaine.meshes.geometry import Plane, xOz_Plane

RNG = np.random.default_rng()


# Some meshes that will be used in the following tests.
test_mesh = OldMesh(vertices=RNG.random((4, 3)), faces=[range(4)], name="test_mesh")
cylinder = mesh_horizontal_cylinder()
sphere = mesh_sphere(radius=1)


def test_mesh_initialization():
    """Test how the code checks the validity of the parameters."""
    with pytest.raises(AssertionError):
        OldMesh(vertices=RNG.random((6, 3)), faces=[(0, 1, 2), (3, 4, 5)])

    with pytest.raises(AssertionError):
        OldMesh(vertices=RNG.random((4, 3)), faces=[(0, 1, 2, -1)])

    with pytest.raises(AssertionError):
        OldMesh(vertices=RNG.random((3, 3)), faces=[(0, 1, 2, 3)])


def test_vertices_and_faces_get_and_set():
    """Test get and set functions for vertices and faces."""
    # Get faces
    faces = cylinder.faces
    vertices = cylinder.vertices

    new_cylinder = OldMesh(name="new_cylinder")

    # Set faces
    with pytest.raises(AssertionError):
        new_cylinder.faces = faces
        # You should set the vertices first.
    new_cylinder.vertices = vertices
    new_cylinder.faces = faces

    assert new_cylinder == cylinder


def test_mesh_naming():
    """Test how the mesh handle names and string representation."""
    # Test automatic naming
    dummy_mesh = OldMesh()  # Automatically named something like mesh_1
    other_dummy_mesh = OldMesh()  # Automatically named something like mesh_2
    assert dummy_mesh.name[:5] == "mesh_"
    assert other_dummy_mesh.name[:5] == "mesh_"
    assert int(dummy_mesh.name[5:]) + 1 == int(other_dummy_mesh.name[5:])


def test_as_set_of_faces():
    """Test the representation of the mesh as a set of faces.
    Also allows to define the equality of two meshes."""
    faces = cylinder.as_set_of_faces()
    assert isinstance(faces, frozenset)
    assert len(faces) == cylinder.nb_faces
    assert all(len(face) == 4 or len(face) == 3 for face in faces)  # Each face is represented by 3 or 4 points
    assert all(len(vertex) == 3 for face in faces for vertex in face)  # Each point is represented by 3 coordinates.

    assert cylinder == cylinder  # The equality is defined as the equality of the set of faces.
    assert OldMesh.from_set_of_faces(faces) == cylinder  # The mesh can be reconstructed from a set of faces.


def mesh_from_set_of_faces():
    A, B, C, D = (0, 0, 0), (1, 0, 0), (0, 1, 0), (1, 1, 0)

    # Two triangular faces
    faces = {frozenset({A, B, C}), frozenset({B, C, D})}
    mesh = cpt.Mesh.from_set_of_faces(faces)
    assert mesh.nb_vertices == 4
    assert mesh.nb_faces == 2
    assert mesh.is_triangle(0) and mesh.is_triangle(1)


def test_copy():
    """Test the copy and renaming of meshes."""
    assert test_mesh.copy() is not test_mesh
    assert test_mesh.copy() == test_mesh
    assert test_mesh.copy(name="copy_of_the_test_mesh").name == "copy_of_the_test_mesh"
    assert test_mesh.merged() is test_mesh


def test_shape_helper_functions():
    assert np.allclose(sphere.center_of_mass_of_nodes, (0.0, 0.0, 0.0))
    assert np.isclose(sphere.diameter_of_nodes, 2.0)
    assert cylinder.squared_axis_aligned_bbox == (-5.0, 5.0, -5.0, 5.0, -5.0, 5.0)


def test_translate():
    assert cylinder.translated_x(0) == cylinder
    assert cylinder.translated_y(0) == cylinder
    assert cylinder.translated_z(0) == cylinder
    assert cylinder.translated([0, 0, 0]) == cylinder

    translated_sphere = sphere.translated([1, 2, 3])
    assert np.allclose(translated_sphere.center_of_mass_of_nodes, (1.0, 2.0, 3.0))
    assert np.isclose(translated_sphere.diameter_of_nodes, sphere.diameter_of_nodes)

def test_rotate():
    assert cylinder.rotated_x(0.0) == cylinder
    assert cylinder.rotated_y(0.0) == cylinder
    assert cylinder.rotated_z(0.0) == cylinder

    new_cylinder = cylinder.rotated_z(np.pi/2)
    assert np.allclose(new_cylinder.center_of_mass_of_nodes, cylinder.center_of_mass_of_nodes)
    assert new_cylinder.squared_axis_aligned_bbox == (-5.0, 5.0, -5.0, 5.0, -5.0, 5.0)


def test_mirror():
    new_cylinder = cylinder.mirrored(Plane())
    cylinder.mirror(Plane())
    assert new_cylinder == cylinder


def test_symmetrized():
    from capytaine.meshes.symmetric import ReflectionSymmetricMesh
    sym = cylinder.merged().symmetrized(xOz_Plane)
    assert isinstance(sym, ReflectionSymmetricMesh)


def test_mesh_quality():
    # TODO: Could be improved...
    cylinder.triangulate_quadrangles()
    cylinder.merge_duplicates(atol=1e-5)
    cylinder.heal_mesh()

def test_heal_mesh_removes_degenerate_panels():
    vertices = np.array([(0.0, 0.0, 0.0), (0.0, 1.0, 0.0), (1.0, 1.0, 0.0), (1.0, 0.0, 0.0)])
    faces = np.array([[0, 1, 2, 3], [1, 2, 2, 1]])
    mesh = OldMesh(vertices, faces)
    assert mesh.nb_faces == 2
    mesh.heal_mesh()
    assert mesh.nb_faces == 1

@pytest.mark.xfail
def test_heal_mesh_with_complicated_connectivities():
    vertices = np.array([
        [-8.00000000e+00,  1.65358984e+00, -4.99999996e-02],
        [-8.00000000e+00,  1.65358984e+00,  5.00000003e-02],
        [-8.00000000e+00,  1.74019238e+00, -9.99999998e-02],
        [-8.00000000e+00,  1.74019238e+00, -1.78037182e-10],
        [-8.00000000e+00,  1.74019238e+00,  1.00000000e-01],
        [-8.00000000e+00,  1.82679492e+00, -5.00000002e-02],
        [-8.00000000e+00,  1.82679492e+00,  4.99999997e-02]
        ])
    faces = np.array([
        [5, 3, 6, 5],
        [5, 2, 3, 5],
        [3, 4, 6, 3],
        [3, 4, 6, 3],
        [2, 0, 3, 2],
        [0, 1, 3, 0],
        [1, 4, 3, 1]
        ])
    mesh = OldMesh(vertices, faces)
    mesh.heal_mesh()

def test_extract_one_face():
    i = 2
    one_face = sphere.extract_one_face(i)
    assert np.all(one_face.faces_centers[0] == sphere.faces_centers[i])


def test_join_meshes_masks():
    a = mesh_sphere(resolution=(2, 3))
    b = mesh_sphere(center=(10, 0, 0), resolution=(3, 3))
    joined, masks = a.join_meshes(b, return_masks=True)
    assert np.all(masks[0] ^ masks[1])
    assert sum(masks[0]) == a.nb_faces
    assert sum(masks[1]) == b.nb_faces


##########################
#  Connected components  #
##########################

def test_connected_components_of_empty_mesh():
    from capytaine.meshes.properties import connected_components
    mesh = OldMesh()
    cc = connected_components(mesh)
    assert len(cc) == 0


def test_connected_components_of_sphere():
    from capytaine.meshes.properties import connected_components
    mesh = mesh_sphere()
    cc = connected_components(mesh)
    assert len(cc) == 1


def test_connected_components_of_two_spheres():
    from capytaine.meshes.properties import connected_components
    mesh = mesh_sphere() + mesh_sphere(center=(5, 0, 0))
    cc = connected_components(mesh)
    assert len(cc) == 2


def test_connected_components_of_open_sphere():
    from capytaine.meshes.properties import connected_components
    mesh = mesh_sphere().immersed_part()
    cc = connected_components(mesh)
    assert len(cc) == 1


def test_connected_components_of_torus():
    from capytaine.io.meshio import load_from_meshio
    from capytaine.meshes.properties import connected_components
    pygmsh = pytest.importorskip("pygmsh")
    with pygmsh.occ.Geometry() as geom:
        geom.add_torus((0, 0, 0), 2.0, 0.5, mesh_size=0.3)
        gmsh_mesh = geom.generate_mesh(dim=2)
    mesh = load_from_meshio(gmsh_mesh).heal_mesh()
    cc = connected_components(mesh)
    assert len(cc) == 1


######################################
#  Connected components at waterline #
######################################

def test_connected_components_at_waterline_of_sphere():
    from capytaine.meshes.properties import connected_components_of_waterline
    mesh = mesh_sphere()
    cc = connected_components_of_waterline(mesh)
    assert len(cc) == 1


def test_connected_components_at_waterline_of_two_spheres():
    from capytaine.meshes.properties import connected_components_of_waterline
    mesh = mesh_sphere() + mesh_sphere(center=(5, 0, 0))
    cc = connected_components_of_waterline(mesh)
    assert len(cc) == 2


def test_connected_components_at_waterline_of_open_sphere():
    from capytaine.meshes.properties import connected_components_of_waterline
    mesh = mesh_sphere().immersed_part()
    cc = connected_components_of_waterline(mesh)
    assert len(cc) == 1


def test_connected_components_at_waterline_of_immersed_sphere():
    from capytaine.meshes.properties import connected_components_of_waterline
    mesh = mesh_sphere(center=(0, 0, -5))
    cc = connected_components_of_waterline(mesh, z=0.0)
    assert len(cc) == 0


def test_connected_components_at_waterline_of_torus():
    from capytaine.io.meshio import load_from_meshio
    from capytaine.meshes.properties import connected_components_of_waterline
    pygmsh = pytest.importorskip("pygmsh")
    with pygmsh.occ.Geometry() as geom:
        geom.add_torus((0, 0, 0), 2.0, 0.5, mesh_size=0.3)
        gmsh_mesh = geom.generate_mesh(dim=2)
    mesh = load_from_meshio(gmsh_mesh).heal_mesh()
    cc = connected_components_of_waterline(mesh)
    assert len(cc) == 2


###################################
#  Connected vertices clustering  #
###################################

@pytest.mark.parametrize("faces", [
    np.array([[1, 2], [1, 3], [2, 3], [4, 5], [5, 6], [7, 8], [3, 9]]),
    np.array([[0, 1, 2, 3], [1, 2, 6, 7]]),
])
def test_vertice_clustering(faces: NDArray[np.integer]):
    """Test the clustering algorithm for connected faces & vertices."""
    # Legacy way to cluster
    vertices_components: set[frozenset[int]] = set()
    for set_of_v_in_face in map(frozenset, faces):
        intersecting_components = [c for c in vertices_components if len(c.intersection(set_of_v_in_face)) > 0]
        if len(intersecting_components) == 0:
            vertices_components.add(set_of_v_in_face)
        else:
            for c in intersecting_components:
                vertices_components.remove(c)
            vertices_components.add(frozenset.union(set_of_v_in_face, *intersecting_components))
    # Vectorized clustering
    vert_groups = clustering(faces)
    assert {frozenset({*group}) for group in vert_groups} == vertices_components
