polyhedron 0.1.0

use std::collections::HashSet;
use std::f32::consts::PI;
use std::fmt::Debug;
use std::ops::Mul;

use linear_isomorphic::{InnerSpace, RealField, VectorSpace};

use super::queue::PQueue;
use crate::verification::verify_correctness;
use crate::*;
use crate::{helpers::approx_normals, MeshStorage, PolyHedron, VertData};

/// Configuration parameters for remeshing.
#[derive(Debug, Clone)]
pub struct RemeshingParameters<S>
where
    S: Debug + Clone + Default + num::Float,
{
    /// No triangle should have an angle smaller than this parameter.
    pub minimum_triangle_angle: S,
    /// Angle at which an edge is considered a feature.
    pub dihedral_angle_tolerance: S,
    /// Number of iterations that remeshing will be applied for.
    pub iterations: usize,
}

impl<S> Default for RemeshingParameters<S>
where
    S: Debug + Clone + Default + num::Float,
{
    fn default() -> Self
    {
        Self {
            minimum_triangle_angle: S::from(2.0 * PI / 9.0).unwrap(),
            dihedral_angle_tolerance: S::from(2.0 * PI / 4.0).unwrap(),
            iterations: 10,
        }
    }
}

/// Uses Kobelt's et. al. Simplified remeshing algorithm to generate new
/// triangles and produce a more equilateral triangulation of the input
/// geometry.
pub fn remeshing<S, L, R, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    parameters: RemeshingParameters<S>,
    adaptive_target_length: L,
    reproject: R,
) where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField
        + Mul<VertData<C>, Output = VertData<C>>
        + ordered_float::FloatCore
        + nalgebra::ComplexField,
    L: Fn(usize, usize) -> S,
    R: Fn(&VertData<C>) -> VertData<C>,
{
    debug_assert!(verify_correctness(mesh).is_ok());

    let (mut feat_verts, mut feat_edges, mut corners) =
        detect_features(mesh, parameters.dihedral_angle_tolerance);

    let low = S::from(4.0 / 5.0).unwrap();
    let high = S::from(4.0 / 3.0).unwrap();

    for _i in 0..parameters.iterations
    {
        split_long_edges(
            mesh,
            high,
            &adaptive_target_length,
            &mut feat_verts,
            &mut feat_edges,
            0,
        );
        debug_assert!(verify_correctness(mesh).is_ok());

        collapse_short_edges(
            mesh,
            low,
            high,
            &adaptive_target_length,
            &mut feat_verts,
            &mut feat_edges,
            &mut corners,
        );

        debug_assert!(verify_correctness(mesh).is_ok());

        equalize_valences(mesh, parameters.minimum_triangle_angle, &feat_edges);
        tangential_relaxation(mesh, &corners);

        debug_assert!(verify_correctness(mesh).is_ok());
        for mut v in mesh.iter_verts()
        {
            let proj = reproject(v.read_data());
            v.update_data(|v| *v = proj.clone());
        }
    }
}

// +| Internal |+ ==============================================================
fn split_long_edges_with_queue<S, L, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    high: S,
    adaptive_target_length: &L,
    feature_verts: &mut HashSet<VertId>,
    feature_edges: &mut HashSet<(VertId, VertId)>,
    pq: &mut PQueue<HalfEdgeId, S>,
    iter_cap: usize,
    return_modified: bool,
) -> (Vec<VertId>, Vec<FaceId>)
where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField
        + Mul<VertData<C>, Output = VertData<C>>
        + ordered_float::FloatCore
        + nalgebra::ComplexField,
    L: Fn(usize, usize) -> S,
{
    let mut mod_verts = Vec::new();
    let mut mod_faces = Vec::new();

    let mut iter_count = 0;
    while !pq.is_empty() && (iter_count < iter_cap || iter_cap == 0)
    {
        let (e_id, weight) = pq.pop().unwrap();
        let hedge_handle = mesh.half_edge_handle(e_id);
        debug_assert!(!hedge_handle.id().is_void());

        let adaptive_length = adaptive_target_length(
            hedge_handle.source().id().0 as usize,
            hedge_handle.dest().id().0 as usize,
        );
        let target_length = high * adaptive_length;
        debug_assert!(target_length > S::from(f32::EPSILON * 10.0).unwrap());
        if S::from(weight).unwrap() > target_length
        {
            let hedge_handle = mesh.half_edge_handle(e_id);
            // TODO(low): properly handle boundary edges in remeshing.
            if hedge_handle.is_in_boundary()
            {
                continue;
            }

            let in_feature = feature_edges
                .contains(&(hedge_handle.source().id(), hedge_handle.dest().id()));

            let v1 = hedge_handle.source().id();
            let v2 = hedge_handle.dest().id();

            let vid = hedge_handle.split::<S>();
            debug_assert!(verify_correctness(mesh).is_ok());

            if return_modified
            {
                mod_verts.push(vid);
                mod_faces.extend(
                    mesh.vert_handle(vid)
                        .iter_star_half_edges()
                        .map(|h| h.face().id()),
                );
            }

            if in_feature
            {
                feature_verts.insert(vid);

                let this_edge = mesh
                    .vert_handle(vid)
                    .iter_star_half_edges()
                    .find(|h| h.dest().id() == v1)
                    .unwrap();
                let other_edge = mesh
                    .vert_handle(vid)
                    .iter_star_half_edges()
                    .find(|h| h.dest().id() == v2)
                    .unwrap();

                let e = mesh.half_edge_handle(this_edge.orbit_next().id());
                feature_edges.insert((e.source().id(), e.dest().id()));
                feature_edges.insert((e.dest().id(), e.source().id()));

                let e = mesh.half_edge_handle(other_edge.orbit_next().id());
                feature_edges.insert((e.source().id(), e.dest().id()));
                feature_edges.insert((e.dest().id(), e.source().id()));
            }

            iter_count += 1;
        }
    }

    (mod_verts, mod_faces)
}

fn split_long_edges<L, S, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    high: S,
    adaptive_target_length: &L,
    feature_verts: &mut HashSet<VertId>,
    feature_edges: &mut HashSet<(VertId, VertId)>,
    iter_cap: usize,
) where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField
        + Mul<VertData<C>, Output = VertData<C>>
        + ordered_float::FloatCore
        + nalgebra::ComplexField,
    L: Fn(usize, usize) -> S,
{
    debug_assert!(verify_correctness(mesh).is_ok());
    let mut pq = PQueue::from_iter(mesh.iter_edges().map(|e| {
        (
            e.id(),
            (e.source().read_data().clone() - e.dest().read_data().clone()).norm(),
        )
    }));

    split_long_edges_with_queue(
        mesh,
        high,
        adaptive_target_length,
        feature_verts,
        feature_edges,
        &mut pq,
        iter_cap,
        false,
    );
}

fn collapse_short_edges<L, S, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    low: S,
    high: S,
    adaptive_target_length: &L,
    feature_verts: &mut HashSet<VertId>,
    feature_edges: &mut HashSet<(VertId, VertId)>,
    corner_verts: &mut HashSet<VertId>,
) where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField
        + Mul<VertData<C>, Output = VertData<C>>
        + ordered_float::FloatCore
        + nalgebra::ComplexField,
    L: Fn(usize, usize) -> S,
{
    let is_feature_vert = |v: &VertHandle<C, MANIFOLD>| feature_verts.contains(&v.id());

    let is_feature_edge = |e: &HalfEdgeHandle<C, MANIFOLD>| {
        feature_edges.contains(&(e.source().id(), e.dest().id()))
    };

    let mut pq = PQueue::from_iter(mesh.iter_edges().map(|e| {
        (
            e.id(),
            (e.source().read_data().clone() - e.dest().read_data().clone()).norm(),
        )
    }));

    debug_assert!(verify_correctness(mesh).is_ok());
    while !pq.is_empty()
    {
        let (e_id, weight) = pq.pop().unwrap();
        let e_handle = mesh.half_edge_handle(e_id);
        if mesh.half_edge_handle(e_id).id().is_void()
        {
            continue;
        }

        let target_length = adaptive_target_length(
            e_handle.source().id().0 as usize,
            e_handle.dest().id().0 as usize,
        );

        if S::from(weight).unwrap() < low * target_length
        {
            // If an edge is not a feature and connects to a feature vert, do not
            // collapse.
            if !is_feature_edge(&e_handle)
                && (is_feature_vert(&e_handle.source())
                    || is_feature_vert(&e_handle.dest()))
            {
                continue;
            }
            // Never collapse boundary hedges.
            if e_handle.is_in_boundary()
            {
                continue;
            }

            // Never collapse edges connecting to the corners.
            if corner_verts.contains(&e_handle.source().id())
                || corner_verts.contains(&e_handle.dest().id())
            {
                continue;
            }

            let a = e_handle.source();
            let b = e_handle.dest();

            let b = b.read_data().clone();

            let mut collapse_ok = true;
            for n in a.iter_neighbours()
            {
                if S::from((b.clone() - n.read_data().clone()).norm()).unwrap()
                    > high * target_length
                {
                    collapse_ok = false;
                }
            }

            if collapse_ok && e_handle.can_collapse()
            {
                let v_id = e_handle.collapse();
                mesh.vert_handle(v_id).update_data(|v| *v = b.clone());
                debug_assert!(verify_correctness(mesh).is_ok());
            }
        }
    }

    debug_assert!(verify_correctness(mesh).is_ok());
}

fn equalize_valences<S, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    minimum_angle: S,
    feature_edges: &HashSet<(VertId, VertId)>,
) where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField + Mul<VertData<C>, Output = VertData<C>>,
{
    let is_feature_edge = |e: &HalfEdgeHandle<C, MANIFOLD>| {
        feature_edges.contains(&(e.source().id(), e.dest().id()))
    };

    let target_val = |v: &VertHandle<C, MANIFOLD>| {
        if v.is_in_boundary()
        {
            4
        }
        else
        {
            6
        }
    };

    let deviation = |hedge: &HalfEdgeHandle<C, MANIFOLD>| {
        let a = hedge.source();
        let b = hedge.dest();
        let c = hedge.face_prev().source();
        let d = hedge.orbit_next().face_prev().source();

        (a.valence() as isize - target_val(&a) as isize).abs()
            + (b.valence() as isize - target_val(&b) as isize).abs()
            + (c.valence() as isize - target_val(&c) as isize).abs()
            + (d.valence() as isize - target_val(&d) as isize).abs()
    };

    debug_assert!(verify_correctness(mesh).is_ok());
    for e in mesh.iter_edges()
    {
        // Don't flip if this would either break topology or flip a feature edge.
        if !e.can_flip() || e.is_in_boundary() || is_feature_edge(&e)
        {
            continue;
        }

        let [o1, o2] = e.opposite_verts();
        let [v1, v2] = [e.source(), e.dest()];

        fn smallest_angle<S: RealField, V: InnerSpace<S>>(v1: V, v2: V, v3: V) -> S
        {
            let d1 = (v2.clone() - v1.clone()).normalized();
            let d2 = (v3.clone() - v2.clone()).normalized();
            let d3 = (v1.clone() - v3.clone()).normalized();

            let a1 = (-d1.dot(&d2)).acos();
            let a2 = (-d2.dot(&d3)).acos();
            let a3 = (-d3.dot(&d1)).acos();

            a1.min(a2).min(a3)
        }

        let post_min_angle = smallest_angle(
            o1.read_data().clone(),
            o2.read_data().clone(),
            v1.read_data().clone(),
        )
        .min(smallest_angle(
            o1.read_data().clone(),
            o2.read_data().clone(),
            v2.read_data().clone(),
        ));

        if S::from(post_min_angle).unwrap() < minimum_angle
        {
            continue;
        }

        let deviation_pre = deviation(&e);
        e.flip();
        let deviation_post = deviation(&e);

        if deviation_pre <= deviation_post
        {
            e.flip();
            debug_assert!(verify_correctness(mesh).is_ok());
        }
    }
}

fn tangential_relaxation<S, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    corner_verts: &HashSet<VertId>,
) where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField + Mul<VertData<C>, Output = VertData<C>>,
{
    let normals = approx_normals(
        mesh.vert_count(),
        mesh.iter_faces().map(|f| {
            f.half_edge()
                .iter_face()
                .map(|e| (e.source().read_data().clone(), e.source().id().to_index()))
        }),
    );

    let mut barycenters = vec![VertData::<C>::default(); mesh.vert_count()];
    let mut valences = vec![0; mesh.vert_count()];
    for vert in mesh.iter_verts()
    {
        let mut valence = 0;
        for neighbour in vert.iter_neighbours()
        {
            barycenters[neighbour.id().to_index()] += neighbour.read_data().clone();
            valence += 1;
        }
        valences[vert.id().to_index()] = valence;
    }

    for (i, &valence) in valences.iter().enumerate()
    {
        barycenters[i] =
            barycenters[i].clone() * (S::from(1.0).unwrap() / S::from(valence).unwrap());
    }

    for mut vert in mesh.iter_verts()
    {
        if corner_verts.contains(&vert.id())
        {
            continue;
        }
        let i = vert.id().to_index();
        let normal = normals[i].clone();
        let mut pos = vert.read_data().clone();
        pos = pos.clone()
            + normal.clone() * normal.dot(&(pos.clone() - barycenters[i].clone()));
        vert.update_data(|v| *v = pos.clone());
    }
}

fn detect_features<S, C: MeshStorage, const MANIFOLD: bool>(
    mesh: &mut PolyHedron<C, MANIFOLD>,
    dihedral_angle_tolerance: S,
) -> (
    HashSet<VertId>,           // Feature verts.
    HashSet<(VertId, VertId)>, // Feature edges.
    HashSet<VertId>,           // Corner verts.
)
where
    VertData<C>: linear_isomorphic::InnerSpace<S>,
    S: linear_isomorphic::RealField + Mul<VertData<C>, Output = VertData<C>>,
{
    let is_feature_edge = |e: &HalfEdgeHandle<C, MANIFOLD>| {
        let [o1, o2] = e.opposite_verts();
        let [v1, v2] = [e.source(), e.dest()];

        let n1 = (v1.read_data().clone() - o1.read_data().clone())
            .cross(&(v2.read_data().clone() - o1.read_data().clone()))
            .normalized();

        let n2 = (v1.read_data().clone() - o2.read_data().clone())
            .cross(&(v2.read_data().clone() - o2.read_data().clone()))
            .normalized();

        let dihedral_angle = n1.dot(&n2).acos();

        S::from(dihedral_angle).unwrap() < dihedral_angle_tolerance
    };

    let mut feature_verts = HashSet::new();
    let mut feature_edges = HashSet::new();
    for edge in mesh.iter_edges()
    {
        if is_feature_edge(&edge)
        {
            feature_verts.insert(edge.source().id());
            feature_verts.insert(edge.dest().id());

            feature_edges.insert((edge.source().id(), edge.dest().id()));
            feature_edges.insert((edge.dest().id(), edge.source().id()));
        }
    }

    // Find corners in the mesh.
    let mut corner_verts = HashSet::new();
    {
        let is_feature_edge = |e: &HalfEdgeHandle<C, MANIFOLD>| {
            feature_edges.contains(&(e.source().id(), e.dest().id()))
        };

        for v in mesh.iter_verts()
        {
            let mut feature_edge_count = 0;
            for e in v.iter_star_half_edges()
            {
                feature_edge_count += is_feature_edge(&e) as usize;
            }

            if feature_edge_count != 2 && feature_edge_count != 0
            {
                corner_verts.insert(v.id());
            }
        }
    }

    (feature_verts, feature_edges, corner_verts)
}
// +| Tests |+ =======================================================
#[cfg(test)]
mod tests
{
    use wavefront_loader::ObjData;

    use super::*;

    #[test]
    fn test_remeshing()
    {
        let ObjData {
            vertices,
            vertex_face_indices,
            ..
        } = ObjData::from_disk_file("../../../Assets/low_poly_bunny.obj");

        let mut mesh = BasicPolyHedron::<_, false>::new(
            vertices.clone(),
            (),
            vertex_face_indices
                .iter()
                .map(|l| l.iter().map(|i| *i as usize)),
        );
        verify_correctness(&mesh).unwrap();

        let length = mesh.half_edge_handle(HalfEdgeId::new(0)).dir().norm() * 0.1;

        let (vs, fs) = mesh.face_list();
        ObjData::export(&(&vs, &fs), "tmp/pre_remeshed_bunny.obj");
        remeshing(
            &mut mesh,
            super::RemeshingParameters {
                minimum_triangle_angle: (20.0_f32).to_radians(),
                dihedral_angle_tolerance: (135.0_f32).to_radians(),
                iterations: 10,
            },
            |_, _| length,
            |a| *a,
        );

        let (vs, fs) = mesh.face_list();
        ObjData::export(&(&vs, &fs), "tmp/remeshed_bunny.obj");
    }
}