mod extended;

pub use extended::*;

use buffer::DataBuffer;
use crate::mesh::attrib::*;
use crate::mesh::topology::*;
use crate::mesh::vertex_positions::VertexPositions;
use crate::prim::Triangle;
use crate::Real;
use reinterpret::*;
#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use std::slice::{Iter, IterMut};
use crate::mesh::PolyMesh;

/*
 * Commonly used meshes and their implementations.
 */

macro_rules! impl_uniform_surface_mesh {
    ($mesh_type:ident, $verts_per_face:expr) => {
        impl<T: Real> $mesh_type<T> {
            pub fn new(verts: Vec<[T; 3]>, indices: Vec<usize>) -> $mesh_type<T> {
                // TODO: Refactor the following unnecessary unsafe block.
                $mesh_type {
                    vertex_positions: IntrinsicAttribute::from_vec(verts),
                    indices: IntrinsicAttribute::from_vec(unsafe { reinterpret_vec(indices) }),
                    vertex_attributes: AttribDict::new(),
                    face_attributes: AttribDict::new(),
                    face_vertex_attributes: AttribDict::new(),
                    face_edge_attributes: AttribDict::new(),
                }
            }

            /// Iterate over each face.
            pub fn face_iter(&self) -> Iter<[usize; $verts_per_face]> {
                self.indices.iter()
            }

            /// Iterate mutably over each face.
            pub fn face_iter_mut(&mut self) -> IterMut<[usize; $verts_per_face]> {
                self.indices.iter_mut()
            }

            /// Face accessor. These are vertex indices.
            #[inline]
            pub fn face(&self, fidx: FaceIndex) -> &[usize; $verts_per_face] {
                &self.indices[fidx]
            }

            /// Return a slice of individual faces.
            #[inline]
            pub fn faces(&self) -> &[[usize; $verts_per_face]] {
                self.indices.as_slice()
            }

            /// Reverse the order of each polygon in this mesh.
            #[inline]
            pub fn reverse(&mut self) {
                for face in self.face_iter_mut() {
                    face.reverse();
                }
            }

            /// Reverse the order of each polygon in this mesh. This is the consuming version of the
            /// `reverse` method.
            #[inline]
            pub fn reversed(mut self) -> Self {
                self.reverse();
                self
            }

            /// Sort vertices by the given key values.
            pub fn sort_vertices_by_key<K, F>(&mut self, f: F)
            where
                F: FnMut(usize) -> K,
                K: Ord,
            {
                if self.num_vertices() == 0 {
                    return;
                }
                self.sort_vertices_by_key_impl(f);
            }

            /// Sort vertices by the given key values, and return the reulting order (permutation).
            /// This function assumes we have at least one vertex.
            pub(crate) fn sort_vertices_by_key_impl<K, F>(&mut self, mut f: F) -> Vec<usize>
            where
                F: FnMut(usize) -> K,
                K: Ord,
            {
                let num = self.attrib_size::<VertexIndex>();
                debug_assert!(num > 0);

                // Original vertex indices.
                let mut order: Vec<usize> = (0..num).collect();

                // Sort vertex indices by the given key.
                order.sort_by_key(|k| f(*k));

                // Now sort all mesh data according to the sorting given by order.

                let $mesh_type {
                    ref mut vertex_positions,
                    ref mut indices,
                    ref mut vertex_attributes,
                    .. // face and face_{vertex,edge} attributes are unchanged
                } = *self;

                let mut seen = vec![false; vertex_positions.len()];

                // Apply the order permutation to vertex_positions in place
                apply_permutation(&order, vertex_positions.as_mut_slice(), &mut seen);

                // Apply permutation to each vertex attribute
                for (_, attrib) in vertex_attributes.iter_mut() {
                    let buf_mut = attrib.buffer_mut();
                    let stride = buf_mut.element_size();
                    let data = buf_mut.as_bytes_mut();

                    apply_permutation_with_stride(&order, data, stride, &mut seen);
                }

                // Build a reverse mapping for convenience.
                let mut new_indices = vec![0; order.len()];
                for (new_idx, &old_idx) in order.iter().enumerate() {
                    new_indices[old_idx] = new_idx;
                }

                // Remap face vertices.
                for face in indices.iter_mut() {
                    for vtx_idx in face.iter_mut() {
                        *vtx_idx = new_indices[*vtx_idx];
                    }
                }

                order
            }
        }

        impl<T: Real> NumVertices for $mesh_type<T> {
            fn num_vertices(&self) -> usize {
                self.vertex_positions.len()
            }
        }

        impl<T: Real> NumFaces for $mesh_type<T> {
            fn num_faces(&self) -> usize {
                self.indices.len()
            }
        }

        impl<T: Real> FaceVertex for $mesh_type<T> {
            #[inline]
            fn face_to_vertex<FI>(&self, fidx: FI, which: usize) -> Option<VertexIndex>
            where
                FI: Copy + Into<FaceIndex>,
            {
                if which >= $verts_per_face {
                    None
                } else {
                    Some(self.indices[fidx.into()][which].into())
                }
            }

            #[inline]
            fn face_vertex<FI>(&self, fidx: FI, which: usize) -> Option<FaceVertexIndex>
            where
                FI: Copy + Into<FaceIndex>,
            {
                if which >= $verts_per_face {
                    None
                } else {
                    let fidx = usize::from(fidx.into());
                    Some(($verts_per_face * fidx + which).into())
                }
            }

            #[inline]
            fn num_face_vertices(&self) -> usize {
                self.indices.len() * $verts_per_face
            }

            #[inline]
            fn num_vertices_at_face<FI>(&self, _: FI) -> usize
            where
                FI: Copy + Into<FaceIndex>,
            {
                $verts_per_face
            }
        }

        impl<T: Real> FaceEdge for $mesh_type<T> {
            #[inline]
            fn face_to_edge<FI>(&self, fidx: FI, which: usize) -> Option<EdgeIndex>
            where
                FI: Copy + Into<FaceIndex>,
            {
                // Edges are assumed to be indexed the same as face vertices: the source of each
                // edge is the face vertex with the same index.
                if which >= $verts_per_face {
                    None
                } else {
                    Some(self.indices[fidx.into()][which].into())
                }
            }
            #[inline]
            fn face_edge<FI>(&self, fidx: FI, which: usize) -> Option<FaceEdgeIndex>
            where
                FI: Copy + Into<FaceIndex>,
            {
                // Edges are assumed to be indexed the same as face vertices: the source of each
                // edge is the face vertex with the same index.
                if which >= $verts_per_face {
                    None
                } else {
                    let fidx = usize::from(fidx.into());
                    Some(($verts_per_face * fidx + which).into())
                }
            }
            #[inline]
            fn num_face_edges(&self) -> usize {
                self.indices.len() * $verts_per_face
            }
            #[inline]
            fn num_edges_at_face<FI>(&self, _: FI) -> usize
            where
                FI: Copy + Into<FaceIndex>,
            {
                $verts_per_face
            }
        }

        impl<T: Real> Default for $mesh_type<T> {
            /// Produce an empty mesh. This is not particularly useful on its own, however it can be
            /// used as a null case for various mesh algorithms.
            fn default() -> Self {
                $mesh_type::new(vec![], vec![])
            }
        }
    };
}

#[derive(Clone, Debug, PartialEq, Attrib, Intrinsic)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct TriMesh<T: Real> {
    /// Vertex positions.
    #[intrinsic(VertexPositions)]
    pub vertex_positions: IntrinsicAttribute<[T; 3], VertexIndex>,
    /// Triples of indices into `vertices` representing triangles.
    pub indices: IntrinsicAttribute<[usize; 3], FaceIndex>,
    /// Vertex attributes.
    pub vertex_attributes: AttribDict<VertexIndex>,
    /// Triangle attributes.
    pub face_attributes: AttribDict<FaceIndex>,
    /// Triangle vertex attributes.
    pub face_vertex_attributes: AttribDict<FaceVertexIndex>,
    /// Triangle edge attributes.
    pub face_edge_attributes: AttribDict<FaceEdgeIndex>,
}

#[derive(Clone, Debug, PartialEq, Attrib, Intrinsic)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct QuadMesh<T: Real> {
    /// Vertex positions.
    #[intrinsic(VertexPositions)]
    pub vertex_positions: IntrinsicAttribute<[T; 3], VertexIndex>,
    /// Quadruples of indices into `vertices` representing quadrilaterals.
    pub indices: IntrinsicAttribute<[usize; 4], FaceIndex>,
    /// Vertex attributes.
    pub vertex_attributes: AttribDict<VertexIndex>,
    /// Quad attributes.
    pub face_attributes: AttribDict<FaceIndex>,
    /// Quad vertex attributes.
    pub face_vertex_attributes: AttribDict<FaceVertexIndex>,
    /// Quad edge attributes.
    pub face_edge_attributes: AttribDict<FaceEdgeIndex>,
}

impl_uniform_surface_mesh!(TriMesh, 3);
impl_uniform_surface_mesh!(QuadMesh, 4);

impl<T: Real> TriMesh<T> {
    /// Triangle iterator.
    ///
    /// ```
    /// use gut::mesh::TriMesh;
    /// use gut::prim::Triangle;
    ///
    /// let verts = vec![[0.0, 0.0, 0.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0]];
    /// let mesh = TriMesh::new(verts.clone(), vec![0, 1, 2]);
    /// let tri = Triangle::from_indexed_slice(&[0, 1, 2], verts.as_slice());
    /// assert_eq!(Some(tri), mesh.tri_iter().next());
    /// ```
    #[inline]
    pub fn tri_iter<'a>(&'a self) -> impl Iterator<Item = Triangle<T>> + 'a {
        self.face_iter().map(move |tri| self.tri_from_indices(tri))
    }

    /// Get a tetrahedron primitive corresponding to the given vertex indices.
    #[inline]
    pub fn tri_from_indices(&self, indices: &[usize; 3]) -> Triangle<T> {
        Triangle::from_indexed_slice(indices, self.vertex_positions.as_slice())
    }
}

/// Convert a triangle mesh to a polygon mesh.
// TODO: Improve this algorithm with ear clipping:
// https://www.geometrictools.com/Documentation/TriangulationByEarClipping.pdf
//
// Ear Clipping Notes:
// 1. Since Polygons are embedded in 3D instead of 2D, we choose to ignore potential intersections
//    of the ears with other ears. This can be resolved as a post process.
//    This means that criteria for bein an ear is only the angle (no vertex inclusion test).
// 2. Polygons have an orientation. We use this (Right-hand-rule) orientation to determine whether
//    a vertex is convex (<180 degrees) or reflex (>180 degrees).
//
// Ear Clipping Algorithm:
// ...TBD
impl<T: Real> From<PolyMesh<T>> for TriMesh<T> {
    fn from(mesh: PolyMesh<T>) -> TriMesh<T> {
        let mut tri_indices = Vec::with_capacity(mesh.num_faces());
        let mut tri_face_attributes: AttribDict<FaceIndex> = AttribDict::new();
        let mut tri_face_vertex_attributes: AttribDict<FaceVertexIndex> = AttribDict::new();
        let mut tri_face_edge_attributes: AttribDict<FaceEdgeIndex> = AttribDict::new();

        // A mapping back to vertices from the polymesh. This allows us to transfer face vertex
        // attributes.
        let mut poly_face_vert_map: Vec<usize> = Vec::with_capacity(mesh.num_face_vertices());

        // Triangulate
        for (face_idx, face) in mesh.face_iter().enumerate() {
            if face.len() < 3 {
                continue;
            }

            let mut idx_iter = face.iter();
            let first_idx = idx_iter.next().unwrap();
            let mut second_idx = idx_iter.next().unwrap();
            let mut second = 1;

            for idx in idx_iter {
                tri_indices.push([*first_idx, *second_idx, *idx]);
                poly_face_vert_map.push(mesh.face_vertex(face_idx, 0).unwrap().into());
                poly_face_vert_map.push(mesh.face_vertex(face_idx, second).unwrap().into());
                second += 1;
                poly_face_vert_map.push(mesh.face_vertex(face_idx, second).unwrap().into());
                second_idx = idx;
            }
        }

        // Transfer face vertex attributes
        for (name, attrib) in mesh.attrib_dict::<FaceVertexIndex>().iter() {
            let mut data = DataBuffer::with_buffer_type(attrib.buffer_ref());
            let attrib_bytes = attrib.buffer_ref();

            for &poly_face_vtx_idx in poly_face_vert_map.iter() {
                data.push_bytes(attrib_bytes.get_bytes(poly_face_vtx_idx));
            }

            tri_face_vertex_attributes.insert(
                name.to_string(),
                Attribute::from_data_buffer(data, attrib.default_bytes()),
            );
        }

        // Transfer face edge attributes
        // We use the face vertex map here because edges have the same topology as face vertices.
        for (name, attrib) in mesh.attrib_dict::<FaceEdgeIndex>().iter() {
            let mut data = DataBuffer::with_buffer_type(attrib.buffer_ref());
            let attrib_bytes = attrib.buffer_ref();

            for &poly_face_edge_idx in poly_face_vert_map.iter() {
                data.push_bytes(attrib_bytes.get_bytes(poly_face_edge_idx));
            }

            tri_face_edge_attributes.insert(
                name.to_string(),
                Attribute::from_data_buffer(data, attrib.default_bytes()),
            );
        }

        // Transfer face attributes
        for (name, attrib) in mesh.attrib_dict::<FaceIndex>().iter() {
            let mut data = DataBuffer::with_buffer_type(attrib.buffer_ref());
            let attrib_chunks = attrib.buffer_ref().byte_chunks();

            // Copy the attribute for every triangle originating from this polygon.
            for (face, bytes) in mesh.face_iter().zip(attrib_chunks) {
                for _ in 2..face.len() {
                    data.push_bytes(bytes);
                }
            }

            tri_face_attributes.insert(
                name.to_string(),
                Attribute::from_data_buffer(data, attrib.default_bytes()),
            );
        }

        let PolyMesh {
            vertex_positions,
            vertex_attributes,
            ..
        } = mesh;

        TriMesh {
            vertex_positions,
            indices: IntrinsicAttribute::from_vec(tri_indices),
            vertex_attributes,
            face_attributes: tri_face_attributes,
            face_vertex_attributes: tri_face_vertex_attributes,
            face_edge_attributes: tri_face_edge_attributes,
        }
    }
}

// Utility functions
fn apply_permutation<T>(permutation: &[usize], array: &mut [T], seen: &mut [bool]) {
    apply_permutation_with_stride(permutation, array, 1, seen);
}

fn apply_permutation_with_stride<T>(
    permutation: &[usize],
    array: &mut [T],
    stride: usize,
    seen: &mut [bool],
) {
    debug_assert_eq!(permutation.len() * stride, array.len());
    debug_assert_eq!(seen.len() * stride, array.len());

    // Clear seen
    seen.iter_mut().for_each(|x| *x = false);

    for unseen_i in 0..seen.len() {
        if seen[unseen_i] {
            continue;
        }

        let mut i = unseen_i;
        loop {
            let idx = permutation[i];
            if seen[idx] {
                break;
            }

            for off in 0..stride {
                array.swap(off + stride * i, off + stride * idx);
            }

            seen[i] = true;
            i = idx;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::index::Index;

    #[test]
    fn mesh_sort() {
        // Sort -> check for inequality -> sort to original -> check for equality.

        let pts = vec![
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [1.0, 1.0, 0.0],
            [1.0, 1.0, 1.0],
        ];
        let indices = vec![0, 1, 2, 1, 3, 2, 0, 2, 4];

        let mut trimesh = TriMesh::new(pts, indices);

        let orig_trimesh = trimesh.clone();

        let values = [3, 2, 1, 4, 0];
        trimesh.sort_vertices_by_key(|k| values[k]);

        assert_ne!(trimesh, orig_trimesh);

        let rev_values = [4, 2, 1, 0, 3];
        trimesh.sort_vertices_by_key(|k| rev_values[k]);

        assert_eq!(trimesh, orig_trimesh);

        // Verify exact values.
        trimesh
            .add_attrib_data::<usize, VertexIndex>("i", vec![0, 1, 2, 3, 4])
            .unwrap();

        trimesh.sort_vertices_by_key(|k| values[k]);

        assert_eq!(
            trimesh.vertex_positions(),
            &[
                [1.0, 1.0, 1.0],
                [0.0, 1.0, 0.0],
                [1.0, 0.0, 0.0],
                [0.0, 0.0, 0.0],
                [1.0, 1.0, 0.0],
            ]
        );

        // `rev_values` actually already corresponds to 0..=4 being sorted by `values`.
        assert_eq!(
            trimesh.attrib_as_slice::<usize, VertexIndex>("i").unwrap(),
            &rev_values[..]
        );
        assert_eq!(
            trimesh.indices.as_slice(),
            &[[3, 2, 1], [2, 4, 1], [3, 1, 0]]
        );
    }

    #[test]
    fn apply_permutation_test() {
        // Checks inner loop in apply_permutation
        let perm = vec![7, 8, 2, 3, 4, 1, 6, 5, 0];
        let mut values = String::from("tightsemi");
        let mut seen = vec![false; 9];
        apply_permutation(&perm, unsafe { values.as_bytes_mut() }, &mut seen);
        assert_eq!(values, "mightiest");

        // Checks the outer loop in apply_permutation
        let perm = vec![7, 8, 4, 3, 2, 1, 6, 5, 0];
        let mut values = String::from("tightsemi");
        let mut seen = vec![false; 9];
        apply_permutation(&perm, unsafe { values.as_bytes_mut() }, &mut seen);
        assert_eq!(values, "mithgiest");

        let mut pts = vec![
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [1.0, 1.0, 0.0],
            [1.0, 1.0, 1.0],
        ];
        seen.resize(5, false);
        let order = [3, 2, 1, 4, 0];
        apply_permutation(&order, &mut pts, &mut seen);
        assert_eq!(
            pts.as_slice(),
            &[
                [1.0, 1.0, 0.0],
                [0.0, 1.0, 0.0],
                [1.0, 0.0, 0.0],
                [1.0, 1.0, 1.0],
                [0.0, 0.0, 0.0],
            ]
        );
    }

    #[test]
    fn two_triangles() {
        let pts = vec![
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [1.0, 1.0, 0.0],
        ];
        let indices = vec![0, 1, 2, 1, 3, 2];

        let trimesh = TriMesh::new(pts, indices);
        assert_eq!(trimesh.num_vertices(), 4);
        assert_eq!(trimesh.num_faces(), 2);
        assert_eq!(trimesh.num_face_vertices(), 6);
        assert_eq!(trimesh.num_face_edges(), 6);

        assert_eq!(Index::from(trimesh.face_to_vertex(1, 1)), 3);
        assert_eq!(Index::from(trimesh.face_to_vertex(0, 2)), 2);
        assert_eq!(Index::from(trimesh.face_edge(1, 0)), 3);

        let mut face_iter = trimesh.face_iter();
        assert_eq!(face_iter.next(), Some(&[0usize, 1, 2]));
        assert_eq!(face_iter.next(), Some(&[1usize, 3, 2]));

    }

    /// Test converting from a `PolyMesh` into a `TriMesh`, which is a non-trivial operation since
    /// it involves trianguating polygons.
    #[test]
    fn from_polymesh() {
        let points = vec![
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [1.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
            [1.0, 0.0, 1.0],
        ];
        let faces = vec![
            3, 0, 1, 2, // first triangle
            4, 0, 1, 5, 4, // quadrilateral
            3, 1, 3, 2, // second triangle
        ];

        let polymesh = crate::mesh::PolyMesh::new(points.clone(), &faces);
        let trimesh = TriMesh::new(points.clone(), vec![0, 1, 2, 0, 1, 5, 0, 5, 4, 1, 3, 2]);
        assert_eq!(trimesh, TriMesh::from(polymesh));
    }
    
    /// Test converting from a `PolyMesh` into a `TriMesh` with attributes.
    #[test]
    fn from_polymesh_with_attrib() -> Result<(), Error> {
        let points = vec![
            [0.0, 0.0, 0.0],
            [1.0, 0.0, 0.0],
            [0.0, 1.0, 0.0],
            [1.0, 1.0, 0.0],
            [0.0, 0.0, 1.0],
            [1.0, 0.0, 1.0],
        ];
        let faces = vec![
            3, 0, 1, 2, // first triangle
            4, 0, 1, 5, 4, // quadrilateral
            3, 1, 3, 2, // second triangle
        ];

        let mut polymesh = crate::mesh::PolyMesh::new(points.clone(), &faces);
        polymesh.add_attrib_data::<u64, VertexIndex>("v", vec![1, 2, 3, 4, 5, 6])?;
        polymesh.add_attrib_data::<u64, FaceIndex>("f", vec![1, 2, 3])?;
        polymesh.add_attrib_data::<u64, FaceVertexIndex>("vf", vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])?;
        polymesh.add_attrib_data::<u64, FaceEdgeIndex>("ve", vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])?;

        let mut trimesh = TriMesh::new(points.clone(), vec![0, 1, 2, 0, 1, 5, 0, 5, 4, 1, 3, 2]);
        trimesh.add_attrib_data::<u64, VertexIndex>("v", vec![1, 2, 3, 4, 5, 6])?;
        trimesh.add_attrib_data::<u64, FaceIndex>("f", vec![1, 2, 2, 3])?;
        trimesh.add_attrib_data::<u64, FaceVertexIndex>("vf", vec![1, 2, 3, 4, 5, 6, 4, 6, 7, 8, 9, 10])?;
        trimesh.add_attrib_data::<u64, FaceEdgeIndex>("ve", vec![1, 2, 3, 4, 5, 6, 4, 6, 7, 8, 9, 10])?;

        assert_eq!(trimesh, TriMesh::from(polymesh));
        Ok(())
    }
}