/*!
 * This module defines the `Merge` trait which defines how objects are merged together.
 * Implementations for common mesh types are also included here.
 */

use math::{Vector3, Vector4};
use crate::mesh::attrib::*;
use crate::mesh::polymesh::PolyMesh;
use crate::mesh::tetmesh::{TetMesh, TetMeshExt};
use crate::mesh::topology::*;
use crate::mesh::uniform_poly_mesh::{QuadMesh, TriMesh, QuadMeshExt, TriMeshExt};
use crate::Real;

/// A trait describing the action of merging mutltiple objects into a single object of the same
/// type. This represents the inverse of splitting a mesh into components of the same type.
pub trait Merge {
    /// Merge another object into self and return the resulting merged object as a mutable
    /// reference.
    fn merge(&mut self, other: Self) -> &mut Self;

    /// Merge a `Vec` of objecs into one of the same type.
    fn merge_vec(vec: Vec<Self>) -> Self
    where
        Self: Default,
    {
        let n = vec.len();
        let mut iter = vec.into_iter();

        // Handle trivial cases first.
        if n == 0 {
            // If no meshes are given on the input, simply return an empty mesh.
            return Default::default();
        }

        // At this point we know there's at least one object. We will merge onto that.
        let mut obj = iter.next().unwrap();

        if n == 1 {
            // If only one mesh is given, just return it.
            return obj;
        }

        for other in iter {
            obj.merge(other);
        }

        obj
    }

    /// In contrast to `merge_vec`, this function takes an immutable reference to a collection of
    /// meshes, and creates a brand new mesh that is a union of all the given meshes.
    fn merge_slice(slice: &[Self]) -> Self
    where
        Self: Clone + Default,
    {
        let vec = slice.to_vec();
        Self::merge_vec(vec)
    }
}

/// Helper function to merge two attribute tables together. `num_elements` must correspond to the
/// number of elements in each attribute inside `dict`. `num_additional_elements` must
/// correspond to the number of elements in each attribute inside `additional_dict`.
fn merge_attribute_dicts<I>(
    dict: &mut AttribDict<I>,
    num_elements: usize,
    additional_dict: AttribDict<I>,
    num_additional_elements: usize,
) {
    for (name, mut other_attrib) in additional_dict.into_iter() {
        // Check if self already has this attribute. If so, we append, otherwise we create a
        // brand new one.
        match dict.entry(name.to_owned()) {
            Entry::Occupied(entry) => {
                let attrib = entry.into_mut();
                if attrib
                    .buffer_mut()
                    .append(other_attrib.buffer_mut())
                    .is_none()
                {
                    // The copy was unsuccessfuly because attributes have different types.
                    // Since we don't actually know the types, we will simply fill our
                    // attribute with a default value.
                    assert_eq!(other_attrib.len(), num_additional_elements);
                    attrib.extend_by(other_attrib.len());
                }
            }
            Entry::Vacant(entry) => {
                other_attrib.extend_by(num_elements);
                other_attrib.rotate_right(num_elements);
                entry.insert(other_attrib);
            }
        }
    }

    // Extend any attributes in dict that weren't present in `additional_dict` to make sure there
    // are the same number of elements in each attribute in `dict` after this function is
    // completed.

    for (_, attrib) in dict.iter_mut() {
        if attrib.len() < num_elements + num_additional_elements {
            attrib.extend_by(num_additional_elements);
        }

        assert_eq!(attrib.len(), num_elements + num_additional_elements);
    }
}

impl<T: Real> Merge for TetMesh<T> {
    /// Attributes with the same name but different types won't be merged.
    fn merge(&mut self, other: Self) -> &mut Self {
        let self_num_vertices = self.num_vertices();
        let other_num_vertices = other.num_vertices();
        let self_num_cells = self.num_cells();
        let other_num_cells = other.num_cells();
        let self_num_cell_vertices = self.num_cell_vertices();
        let other_num_cell_vertices = other.num_cell_vertices();
        let self_num_cell_faces = self.num_cell_faces();
        let other_num_cell_faces = other.num_cell_faces();

        // Deconstruct the other mesh explicitly since it will not be valid as soon as we start to
        // canibalize its contents.
        let TetMesh {
            vertex_positions: mut other_vertex_positions,
            indices: other_indices,
            vertex_attributes: other_vertex_attributes,
            cell_attributes: other_cell_attributes,
            cell_vertex_attributes: other_cell_vertex_attributes,
            cell_face_attributes: other_cell_face_attributes,
        } = other;

        self.vertex_positions
            .as_mut_vec()
            .append(other_vertex_positions.as_mut_vec());
        self.indices.as_mut_vec().extend(
            other_indices
                .iter()
                .map(|&cell| -> [usize; 4] {
                    Vector4::from(cell).map(|i| i + self_num_vertices).into()
                }),
        );

        // Transfer attributes
        merge_attribute_dicts(
            &mut self.vertex_attributes,
            self_num_vertices,
            other_vertex_attributes,
            other_num_vertices,
        );
        merge_attribute_dicts(
            &mut self.cell_attributes,
            self_num_cells,
            other_cell_attributes,
            other_num_cells,
        );
        merge_attribute_dicts(
            &mut self.cell_vertex_attributes,
            self_num_cell_vertices,
            other_cell_vertex_attributes,
            other_num_cell_vertices,
        );
        merge_attribute_dicts(
            &mut self.cell_face_attributes,
            self_num_cell_faces,
            other_cell_face_attributes,
            other_num_cell_faces,
        );
        self
    }
}

impl<T: Real> Merge for TetMeshExt<T> {
    /// Attributes with the same name but different types won't be merged.
    fn merge(&mut self, other: Self) -> &mut Self {
        let self_num_cells = self.num_cells();
        let self_num_vertex_cells = self.num_vertex_cells();
        let other_num_vertex_cells = other.num_vertex_cells();

        // Deconstruct the other mesh explicitly since it will not be valid as soon as we start to
        // canibalize its contents.
        let TetMeshExt {
            tetmesh: other_tetmesh,
            cell_offsets: other_cell_offsets,
            cell_indices: other_cell_indices,
            vertex_cell_attributes: other_vertex_cell_attributes,
        } = other;

        self.tetmesh.merge(other_tetmesh);
        self.cell_offsets.extend(
            other_cell_offsets
                .iter()
                .skip(1)
                .map(|&i| i + self_num_vertex_cells),
        );
        self.cell_indices
            .extend(other_cell_indices.iter().map(|&i| i + self_num_cells));

        // Transfer attributes
        merge_attribute_dicts(
            &mut self.vertex_cell_attributes,
            self_num_vertex_cells,
            other_vertex_cell_attributes,
            other_num_vertex_cells,
        );
        self
    }
}

impl<T: Real> Merge for PolyMesh<T> {
    /// Attributes with the same name but different types won't be merged.
    fn merge(&mut self, other: Self) -> &mut Self {
        let self_num_vertices = self.num_vertices();
        let other_num_vertices = other.num_vertices();
        let self_num_faces = self.num_faces();
        let other_num_faces = other.num_faces();
        let self_num_face_vertices = self.num_face_vertices();
        let other_num_face_vertices = other.num_face_vertices();
        let self_num_face_edges = self.num_face_edges();
        let other_num_face_edges = other.num_face_edges();

        // Deconstruct the other mesh explicitly since it will not be valid as soon as we start to
        // canibalize its contents.
        let PolyMesh {
            vertex_positions: mut other_vertex_positions,
            indices: other_indices,
            offsets: other_offsets,
            vertex_attributes: other_vertex_attributes,
            face_attributes: other_face_attributes,
            face_vertex_attributes: other_face_vertex_attributes,
            face_edge_attributes: other_face_edge_attributes,
        } = other;

        self.vertex_positions
            .as_mut_vec()
            .append(other_vertex_positions.as_mut_vec());
        self.offsets.extend(
            other_offsets
                .iter()
                .skip(1)
                .map(|&i| i + self_num_face_vertices),
        );
        self.indices
            .extend(other_indices.iter().map(|&i| i + self_num_vertices));

        // Transfer attributes
        merge_attribute_dicts(
            &mut self.vertex_attributes,
            self_num_vertices,
            other_vertex_attributes,
            other_num_vertices,
        );
        merge_attribute_dicts(
            &mut self.face_attributes,
            self_num_faces,
            other_face_attributes,
            other_num_faces,
        );
        merge_attribute_dicts(
            &mut self.face_vertex_attributes,
            self_num_face_vertices,
            other_face_vertex_attributes,
            other_num_face_vertices,
        );
        merge_attribute_dicts(
            &mut self.face_edge_attributes,
            self_num_face_edges,
            other_face_edge_attributes,
            other_num_face_edges,
        );
        self
    }
}

macro_rules! impl_merge_for_uniform_mesh {
    ($mesh_type:ident, $base_type:ident, $verts_per_face:expr, $vec:ident) => {
        impl<T: Real> Merge for $base_type<T> {
            /// Attributes with the same name but different types won't be merged.
            fn merge(&mut self, other: Self) -> &mut Self {
                let self_num_vertices = self.num_vertices();
                let other_num_vertices = other.num_vertices();
                let self_num_faces = self.num_faces();
                let other_num_faces = other.num_faces();
                let self_num_face_vertices = self.num_face_vertices();
                let other_num_face_vertices = other.num_face_vertices();
                let self_num_face_edges = self.num_face_edges();
                let other_num_face_edges = other.num_face_edges();

                // Deconstruct the other mesh explicitly since it will not be valid as soon as we start to
                // canibalize its contents.
                let $base_type {
                    vertex_positions: mut other_vertex_positions,
                    indices: other_indices,
                    vertex_attributes: other_vertex_attributes,
                    face_attributes: other_face_attributes,
                    face_vertex_attributes: other_face_vertex_attributes,
                    face_edge_attributes: other_face_edge_attributes,
                } = other;

                self.vertex_positions
                    .as_mut_vec()
                    .append(other_vertex_positions.as_mut_vec());
                self.indices.as_mut_vec().extend(
                    other_indices
                        .iter()
                        .map(|&face| -> [usize; $verts_per_face] {
                            $vec::from(face).map(|i| i + self_num_vertices).into()
                        }),
                );

                // Transfer attributes
                merge_attribute_dicts(
                    &mut self.vertex_attributes,
                    self_num_vertices,
                    other_vertex_attributes,
                    other_num_vertices,
                );
                merge_attribute_dicts(
                    &mut self.face_attributes,
                    self_num_faces,
                    other_face_attributes,
                    other_num_faces,
                );
                merge_attribute_dicts(
                    &mut self.face_vertex_attributes,
                    self_num_face_vertices,
                    other_face_vertex_attributes,
                    other_num_face_vertices,
                );
                merge_attribute_dicts(
                    &mut self.face_edge_attributes,
                    self_num_face_edges,
                    other_face_edge_attributes,
                    other_num_face_edges,
                );
                self
            }
        }

        impl<T: Real> Merge for $mesh_type<T> {
            /// Attributes with the same name but different types won't be merged.
            fn merge(&mut self, other: Self) -> &mut Self {
                let self_num_faces = self.num_faces();
                let self_num_vertex_faces = self.num_vertex_faces();

                // Deconstruct the other mesh explicitly since it will not be valid as soon as we start to
                // canibalize its contents.
                let $mesh_type {
                    base_mesh: other_base_mesh,
                    face_offsets: other_face_offsets,
                    face_indices: other_face_indices,
                } = other;

                self.base_mesh.merge(other_base_mesh);

                self.face_offsets.extend(
                    other_face_offsets
                        .iter()
                        .skip(1)
                        .map(|&i| i + self_num_vertex_faces),
                );
                self.face_indices
                    .extend(other_face_indices.iter().map(|&i| i + self_num_faces));

                //merge_attribute_dicts(
                //    &mut self.vertex_face_attributes,
                //    self_num_vertex_faces,
                //    other_vertex_face_attributes,
                //    other_num_vertex_faces,
                //    );
                self
            }
        }
    };
}

impl_merge_for_uniform_mesh!(TriMeshExt, TriMesh, 3, Vector3);
impl_merge_for_uniform_mesh!(QuadMeshExt, QuadMesh, 4, Vector4);

#[cfg(test)]
mod tests {
    use super::*;
    use crate::algo::test_utils::*;
    use crate::mesh::TetMeshExt;

    fn build_tetmesh_sample() -> (TetMeshExt<f64>, TetMeshExt<f64>, TetMeshExt<f64>) {
        let verts = vec![
            [0.0, 0.0, 0.0],
            [0.0, 0.0, 1.0],
            [0.0, 1.0, 1.0],
            [0.5, 0.0, 0.5],
            [0.0, 1.0, 0.0],
            [1.0, 0.0, 0.0],
            [1.0, 0.0, 1.0],
            [1.0, 1.0, 0.0],
            [1.0, 1.0, 1.0],
        ];

        // One connected component consisting of two tets connected at a face, and another
        // consisting of a single tet.
        let indices = vec![0, 1, 2, 3, 8, 7, 4, 5, 6, 8, 4, 5];

        let tetmesh = TetMeshExt::new(verts, indices);
        let comp1 = TetMeshExt::new(
            vec![
                [0.0, 0.0, 0.0],
                [0.0, 0.0, 1.0],
                [0.0, 1.0, 1.0],
                [0.5, 0.0, 0.5],
            ],
            vec![0, 1, 2, 3],
        );
        let comp2 = TetMeshExt::new(
            vec![
                [0.0, 1.0, 0.0],
                [1.0, 0.0, 0.0],
                [1.0, 0.0, 1.0],
                [1.0, 1.0, 0.0],
                [1.0, 1.0, 1.0],
            ],
            vec![4, 3, 0, 1, 2, 4, 0, 1],
        );
        (tetmesh, comp1, comp2)
    }

    fn add_attribs_to_tetmeshes(sample: &mut (TetMeshExt<f64>, TetMeshExt<f64>, TetMeshExt<f64>)) {
        // Add a sample vertex attribute.
        sample
            .0
            .add_attrib_data::<usize, VertexIndex>("v", (0..sample.0.num_vertices()).collect())
            .unwrap();
        sample
            .1
            .add_attrib_data::<usize, VertexIndex>("v", (0..4).collect())
            .unwrap();
        sample
            .2
            .add_attrib_data::<usize, VertexIndex>("v", (4..9).collect())
            .unwrap();

        // Add a sample cell attribute.
        sample
            .0
            .add_attrib_data::<usize, CellIndex>("c", (0..sample.0.num_cells()).collect())
            .unwrap();
        sample
            .1
            .add_attrib_data::<usize, CellIndex>("c", vec![0])
            .unwrap();
        sample
            .2
            .add_attrib_data::<usize, CellIndex>("c", vec![1, 2])
            .unwrap();

        // Add a sample cell vertex attribute.
        sample
            .0
            .add_attrib_data::<usize, CellVertexIndex>(
                "cv",
                (0..sample.0.num_cells() * 4).collect(),
            )
            .unwrap();
        sample
            .1
            .add_attrib_data::<usize, CellVertexIndex>("cv", (0..4).collect())
            .unwrap();
        sample
            .2
            .add_attrib_data::<usize, CellVertexIndex>("cv", (4..12).collect())
            .unwrap();

        // Add a sample cell face attribute.
        sample
            .0
            .add_attrib_data::<usize, CellFaceIndex>("cf", (0..sample.0.num_cells() * 4).collect())
            .unwrap();
        sample
            .1
            .add_attrib_data::<usize, CellFaceIndex>("cf", (0..4).collect())
            .unwrap();
        sample
            .2
            .add_attrib_data::<usize, CellFaceIndex>("cf", (4..12).collect())
            .unwrap();

        // Add a sample vertex cell attribute.
        sample
            .0
            .add_attrib_data::<usize, VertexCellIndex>(
                "vc",
                (0..sample.0.num_cells() * 4).collect(),
            )
            .unwrap();
        sample
            .1
            .add_attrib_data::<usize, VertexCellIndex>("vc", (0..4).collect())
            .unwrap();
        sample
            .2
            .add_attrib_data::<usize, VertexCellIndex>("vc", (4..12).collect())
            .unwrap();
    }

    #[test]
    fn tetmesh_merge() {
        // Generate sample meshes
        let mut sample = build_tetmesh_sample();

        // Add all the attributes
        add_attribs_to_tetmeshes(&mut sample);
        let (tetmesh, mut comp1, comp2) = sample;

        comp1.merge(comp2);
        assert_eq!(comp1, tetmesh);
    }

    #[test]
    fn tetmesh_merge_collection() {
        // Generate sample meshes
        let mut sample = build_tetmesh_sample();
        add_attribs_to_tetmeshes(&mut sample);
        let (mesh, comp1, comp2) = sample;

        let res = Merge::merge_vec(vec![comp1.clone(), comp2.clone()]);
        assert_eq!(res, mesh);

        let res = Merge::merge_slice(&[comp1, comp2]);
        assert_eq!(res, mesh);
    }

    #[test]
    fn polymesh_merge() {
        let mut sample = build_polymesh_sample();
        add_attribs_to_polymeshes(&mut sample);
        let (mesh, mut comp1, comp2) = sample;
        comp1.merge(comp2);
        assert_eq!(comp1, mesh);
    }

    #[test]
    fn polymesh_merge_collection() {
        let mut sample = build_polymesh_sample();
        add_attribs_to_polymeshes(&mut sample);
        let (mesh, comp1, comp2) = sample;
        let res = Merge::merge_vec(vec![comp1.clone(), comp2.clone()]);
        assert_eq!(res, mesh);

        let res = Merge::merge_slice(&[comp1, comp2]);
        assert_eq!(res, mesh);
    }

    #[test]
    fn trival_merges() {
        let sample = build_tetmesh_sample();

        // Merging with an empty mesh should return the same mesh back.
        assert_eq!(*TetMeshExt::default().merge(sample.0.clone()), sample.0);
        assert_eq!(*sample.0.clone().merge(TetMeshExt::default()), sample.0);

        // Do the same test with a polymesh
        let sample = build_polymesh_sample();

        assert_eq!(*PolyMesh::default().merge(sample.0.clone()), sample.0);
        assert_eq!(*sample.0.clone().merge(PolyMesh::default()), sample.0);
    }
}