mesh-graph 0.6.0

use hashbrown::HashSet;
use tracing::instrument;

use crate::{
    HalfedgeId, MeshGraph, Selection, SelectionOps, VertexId, error_none, utils::unwrap_or_return,
};

#[cfg(feature = "rerun")]
use crate::utils::vec3_array;

impl MeshGraph {
    /// Subdivide all edges until all of them are <= max_length.
    /// Please note that you have to provide the squared value of max_length.
    ///
    /// This will schedule necessary updates to the QBVH but you have to call
    /// `refit_bvh()` after the operation.
    #[instrument(skip(self))]
    pub fn subdivide_until_edges_below_max_length(
        &mut self,
        max_length_squared: f32,
        marked_halfedge_ids: &mut HashSet<HalfedgeId>,
        marked_vertex_ids: &mut HashSet<VertexId>,
    ) {
        let mut halfedges_to_subdivide = self.halfedges_map(|len_sqr| len_sqr > max_length_squared);

        while !halfedges_to_subdivide.is_empty() {
            let mut max_len = 0.0;
            let mut max_he_id = HalfedgeId::default();

            #[cfg(feature = "rerun")]
            self.log_hes_rerun(
                "subdivide/selection",
                &halfedges_to_subdivide
                    .iter()
                    .map(|(he, _)| *he)
                    .collect::<Vec<_>>(),
            );

            for (&he, &len) in &halfedges_to_subdivide {
                if len > max_len {
                    max_len = len;
                    max_he_id = he;
                }
            }

            halfedges_to_subdivide.remove(&max_he_id);

            let mut affected_faces = Selection::default();

            // already checked
            let max_he = self.halfedges[max_he_id];
            if let Some(face_id) = max_he.face {
                affected_faces.insert(face_id);
            }

            if let Some(twin_id) = max_he.twin.or_else(error_none!("Twin missing"))
                && let Some(twin_he) = self
                    .halfedges
                    .get(twin_id)
                    .or_else(error_none!("Halfedge not found"))
                && let Some(twin_face_id) = twin_he.face
            {
                affected_faces.insert(twin_face_id);
            }

            let subdivide_edge_result =
                unwrap_or_return!(self.subdivide_edge(max_he_id), "Couldn't subdivide edge");

            if marked_halfedge_ids.contains(&max_he_id) {
                marked_halfedge_ids.extend(subdivide_edge_result.added_halfedges.iter().copied());
                marked_vertex_ids.insert(subdivide_edge_result.added_vertex);
            }

            // #[cfg(feature = "rerun")]
            // {
            //     crate::RR
            //         .log("meshgraph/subdivide", &rerun::Clear::recursive())
            //         .unwrap();
            //     crate::RR
            //         .log("meshgraph/halfedge/subdivide", &rerun::Clear::recursive())
            //         .unwrap();
            //     crate::RR
            //         .log("meshgraph/face/subdivide", &rerun::Clear::recursive())
            //         .unwrap();
            //     self.log_rerun();
            // }

            for new_he_id in subdivide_edge_result.added_halfedges {
                // newly inserted in `self.subdivide_edge`
                let new_he = self.halfedges[new_he_id];

                if let Some(face_id) = new_he.face {
                    affected_faces.insert(face_id);
                }

                let new_twin_id = unwrap_or_return!(new_he.twin, "New twin missing");
                let new_twin =
                    unwrap_or_return!(self.halfedges.get(new_twin_id), "New twin not found");

                if let Some(face_id) = new_twin.face {
                    affected_faces.insert(face_id);
                    #[cfg(feature = "rerun")]
                    self.log_face_rerun("subdivide/selected_new", face_id);
                }
            }

            let mut new_hes_to_check = HashSet::new();

            for he_id in affected_faces.resolve_to_halfedges(self) {
                let he = unwrap_or_return!(self.halfedges.get(he_id), "Halfedge not found");
                let twin_id = unwrap_or_return!(he.twin, "Twin missing");

                new_hes_to_check.insert(he_id.min(twin_id));
            }

            for he_id in new_hes_to_check {
                let he = self.halfedges[he_id]; // checked above

                let len_sqr = he.length_squared(self);

                if len_sqr > max_length_squared {
                    #[cfg(feature = "rerun")]
                    self.log_he_rerun("/subdivide/new_edge", he_id);

                    halfedges_to_subdivide.insert(he_id, len_sqr);
                }
            }
        }

        #[cfg(feature = "rerun")]
        self.log_rerun();
    }

    /// Subdivides an edge by computing it's center vertex. This also subdivides any adjacent triangles and
    /// makes sure everything is properly reconnected. Works only on triangle meshes.
    ///
    /// Returns the id of the new halfedge which goes from the center vertex to the original edge's end vertex.
    /// And also return the halfedges that are created by subdividing the adjacent faces. Only one of the two twin
    /// halfedges per face subdivision is returned. In total the number `n` of halfedges returned is `1 <= n <= 3`.
    /// (The one from dividing the halfedge and at most 2 from dividing the two adjacent faces).
    ///
    /// Also returns the created vertex id.
    #[instrument(skip(self))]
    pub fn subdivide_edge(&mut self, halfedge_id: HalfedgeId) -> Option<SubdivideEdge> {
        let mut added_halfedges = Vec::with_capacity(3);

        let he = self
            .halfedges
            .get(halfedge_id)
            .or_else(error_none!("Halfedge not found"))?;
        let twin_id = he.twin.or_else(error_none!("Twin halfedge not found"))?;

        let start_v = he
            .start_vertex(self)
            .or_else(error_none!("Start vertex not found"))?;
        let end_v = he.end_vertex;

        let start_pos = self.positions[start_v];
        let end_pos = self.positions[end_v];

        let center_pos = (start_pos + end_pos) * 0.5;

        #[cfg(feature = "rerun")]
        {
            crate::RR
                .log(
                    "meshgraph/subdivide/edge",
                    &rerun::Arrows3D::from_vectors([vec3_array(end_pos - start_pos)])
                        .with_origins([vec3_array(start_pos)]),
                )
                .unwrap();

            crate::RR
                .log(
                    "meshgraph/subdivide/center",
                    &rerun::Points3D::new([vec3_array(center_pos)]),
                )
                .unwrap();
        }

        let center_v = self.add_vertex(center_pos);
        if let Some(normals) = &mut self.vertex_normals {
            let start_normal = normals
                .get(start_v)
                .or_else(error_none!("Start normal not found"))?;
            let end_normal = normals
                .get(end_v)
                .or_else(error_none!("End normal not found"))?;
            normals.insert(center_v, (start_normal + end_normal).normalize());
        }

        let new_he = self.add_halfedge(center_v, end_v);
        // inserted just above
        self.vertices[center_v].outgoing_halfedge = Some(new_he);

        added_halfedges.push(new_he);

        if let Some(new_face_he) = self.subdivide_face(halfedge_id, new_he, center_v) {
            added_halfedges.push(new_face_he);
        }

        let new_twin = self.add_halfedge(center_v, start_v);

        if let Some(new_face_he) = self.subdivide_face(twin_id, new_twin, center_v) {
            added_halfedges.push(new_face_he);
        }

        // inserted above
        self.halfedges[new_he].twin = Some(twin_id);
        self.halfedges
            .get_mut(twin_id)
            .or_else(error_none!("Twin halfedge not found"))?
            .twin = Some(new_he);

        // checked in the beginning of the function
        self.halfedges[halfedge_id].twin = Some(new_twin);
        // inserted above
        self.halfedges[new_twin].twin = Some(halfedge_id);

        // self.vertices[end_v].outgoing_halfedge = Some(new_twin);
        // self.vertices[start_v].outgoing_halfedge = Some(new_he);

        Some(SubdivideEdge {
            added_halfedges,
            added_vertex: center_v,
        })
    }

    /// Subdivides a triangle into two halves. Used in [Self::subdivide_edge].
    #[instrument(skip(self))]
    fn subdivide_face(
        &mut self,
        existing_halfedge_id: HalfedgeId,
        new_halfedge_id: HalfedgeId,
        center_v: VertexId,
    ) -> Option<HalfedgeId> {
        let he = self
            .halfedges
            .get(existing_halfedge_id)
            .or_else(error_none!("Halfedge not found"))?;

        let face_id = he.face?;
        self.faces
            .get_mut(face_id)
            .or_else(error_none!("Facee not found"))?
            .halfedge = existing_halfedge_id;

        // checked above
        let next_he = self.halfedges[existing_halfedge_id]
            .next
            .or_else(error_none!("Next halfedge missing"))?;
        let last_he = self
            .halfedges
            .get(next_he)
            .or_else(error_none!("Next halfedge not found"))?
            .next
            .or_else(error_none!("Last halfedge not found"))?;

        // rewire existing face
        let new_he = self.add_halfedge(center_v, self.halfedges[next_he].end_vertex);

        self.halfedges[existing_halfedge_id].next = Some(new_he);
        self.halfedges[new_he].next = Some(last_he);
        self.halfedges[new_he].face = Some(face_id);

        let new_twin = self.add_halfedge(self.halfedges[next_he].end_vertex, center_v);

        // insert new face
        let new_face_id = self.add_face(new_halfedge_id, next_he, new_twin);

        self.halfedges[new_twin].twin = Some(new_he);
        self.halfedges[new_he].twin = Some(new_twin);

        self.halfedges[existing_halfedge_id].end_vertex = center_v;

        // checked above
        let face = self.faces[face_id];
        // freshly inserted
        let new_face = self.faces[new_face_id];
        self.bvh
            .insert_or_update_partially(face.aabb(self), face.index, 0.0);
        self.bvh
            .insert_or_update_partially(new_face.aabb(self), new_face.index, 0.0);

        #[cfg(feature = "rerun")]
        {
            self.log_he_rerun("subdivide/new_he", new_he);
            self.log_he_rerun("subdivide/new_twin", new_twin);
        }

        Some(new_he)
    }
}

pub struct SubdivideEdge {
    /// All halfedges created by the subdivision.
    added_halfedges: Vec<HalfedgeId>,
    /// This is the center vertex of the subdivided edge that was created.
    added_vertex: VertexId,
}

#[cfg(test)]
mod tests {}