dyntri-core 0.10.2

use std::collections::{HashMap, VecDeque};
use std::fmt::Debug;
use std::hash::Hash;

use itertools::Itertools;
use log::warn;
use serde::{Deserialize, Serialize};

use crate::graph::Graph;
use crate::triangulation::boundary_map::{
    CanonicalSimplex, Face, FaceSimplex, get_face, get_faces,
};
use crate::triangulation::simplices::{Simplex, canonical_simplex};

/// (N-1)-dimensional triangulation structure designed to store a "gluing" of (N-1)-simplices.
///
/// Internally four half-(N-2)-simplices/"half-faces" (each side of a face gets its own label) are
/// stored with a pointer to the half-face they are adjacent to. Additionally this stores the
/// vertex labels of the N vertices of each (N-1)-simplex opposite the half-faces.
#[derive(Debug, Default, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct Triangulation<const N: usize> {
    /// Number of vertices in the triangulation
    pub num_vertices: usize,
    /// List holding (N-1)-simplices of an (N-1)-dimensional triangulation.
    ///
    /// Note that the `neighbours` field of each 3-simplex is used to store the indices of
    /// the half-face each half-face is adjacent to. This encodes both the label of the
    /// neighbouring (N-1)-simplex and the back index (the index in the `neighbour` list back).
    ///
    /// The encoding works as follows, given the "half-face" label `i`:
    /// - `simplex_label = i / N`
    /// - `face_index = i % N`
    pub simplices: Vec<Simplex<N>>,
}

pub type Triangulation2D = Triangulation<3>;
pub type Triangulation3D = Triangulation<4>;
pub type Triangulation4D = Triangulation<5>;

impl<const N: usize> Triangulation<N> {
    /// Creates and empty triangulation
    pub fn empty() -> Self {
        Self::default()
    }
}

impl<const N: usize> Triangulation<N> {
    /// Returns the number of vertices in the triangulation
    #[inline]
    pub fn num_vertices(&self) -> usize {
        self.num_vertices
    }

    /// Returns the number of (N-1)-simplices in the triangulation
    #[inline]
    pub fn num_simplices(&self) -> usize {
        self.simplices.len()
    }

    /// Returns the number of (N-2)-simplices/faces codim1 in the triangulation
    ///
    /// Note: this is simply `num_simplices * N / 2`
    #[inline]
    pub fn num_faces(&self) -> usize {
        // Division is exact for a valid triangulation
        self.num_halffaces() / 2
    }

    /// Returns the number of half-faces of the triangulation
    ///
    /// Note: this is simply `num_simplices * N`
    #[inline]
    pub fn num_halffaces(&self) -> usize {
        self.simplices.len() * N
    }

    /// Returns the label of the adjacent `halfface`
    #[inline]
    pub fn adjacent_halfface(&self, halfface: usize) -> usize {
        self.simplices[halfface / N].neighbours[halfface % N]
    }

    /// Returns the label of the simplex the `halfface` is part of.
    ///
    /// Note: this is simply `halfface / N`
    #[inline]
    pub fn simplex_label(halfface: usize) -> usize {
        halfface / N
    }

    /// Returns the index of the half-face in the corresponding triangulation.
    ///
    /// This is simply `halfface % N` and means the index of neighbouring (N-1)-simplex
    /// in the `neighbours` list of [`Simplex`] corresponding to the half-face.
    #[inline]
    pub fn halfface_index(halfface: usize) -> usize {
        halfface % N
    }

    /// Returns the label of the half-face given by a `simplex_label` and `halfface_index`
    ///
    /// Note that `halfface_index < N`, for larger indices this function will give a result,
    /// but it does not correspond to a half-edge of the given (N-1)-simplex.
    /// Note this function is simply `N * simplex_label + halfface_index`
    #[inline]
    pub fn halfface_label(simplex_label: usize, halfface_index: usize) -> usize {
        5 * simplex_label + halfface_index
    }
}

impl<const N: usize> Triangulation<N> {
    /// Returns a reference to the (N-1)-simplex list of the [`Triangulation`]
    #[inline]
    pub fn get_simplices(&self) -> &[Simplex<N>] {
        &self.simplices
    }

    /// Returns a reference to the (N-1)-simplex given by `label`
    ///
    /// Note that the `neighbours` field hold the labels of the adjacent half-faces, encoding
    /// the triangle label as `i / N` and the back-index as `i % N`. Also see [`Triangulation`].
    ///
    /// # Panic
    /// This method will panic if `label >= [self.num_simplices()]`
    #[inline]
    pub fn get_simplex(&self, label: usize) -> &Simplex<N> {
        &self.simplices[label]
    }

    /// Returns a mutable reference to the (N-1)-simplex given by `label`
    ///
    /// See [`Triangulation::get_simplex()`] for more details.
    ///
    /// # Panic
    /// This method will panic if `label >= self.num_simplices()`
    #[inline]
    pub fn get_mut_simplex(&mut self, label: usize) -> &mut Simplex<N> {
        &mut self.simplices[label]
    }

    /// Returns the label of the (N-1)-simplex neighbouring `label` on the `index` side.
    ///
    /// # Panic
    /// This method will panic if `label >= self.num_simplices()` or `index >= N`
    #[inline]
    pub fn get_neighbour(&self, label: usize, index: usize) -> usize {
        self.get_simplex(label).get_neighbour(index)
    }

    /// Returns the label and backindex of (N-1)-simplex neighbouring `label` on the `index` side.
    ///
    /// # Panic
    /// This method will panic if `label >= self.num_simplices()` or `index >= N`
    #[inline]
    pub fn get_neighbour_backindex(&self, label: usize, index: usize) -> (usize, usize) {
        self.get_simplex(label).get_neighbour_backindex(index)
    }

    /// Returns the neighbouring (N-1)-simplex labels of the (N-1)-simplex given by `label`
    ///
    /// # Panic
    /// This method will panic if `label >= self.num_simplices()`
    #[inline]
    pub fn get_neighbours(&self, label: usize) -> [usize; N] {
        self.get_simplex(label).get_neighbours()
    }

    /// Returns the neighbouring (N-1)-simplex labels and backindices of the (N-1)-simplex
    /// given by `label`.
    ///
    /// # Panic
    /// This method will panic if `label >= self.num_simplices()`
    #[inline]
    pub fn get_neighbours_backindices(&self, label: usize) -> [(usize, usize); N] {
        self.get_simplex(label).get_neighbours_backindices()
    }

    /// Returns the vertices of the (N-1)-simplex by `label`
    ///
    /// # Panic
    /// This method will panic if `label >= self.num_simplices()`
    #[inline]
    pub fn get_simplex_vertices(&self, label: usize) -> [usize; N] {
        self.get_simplex(label).vertices
    }
}

impl<const N: usize> Triangulation<N> {
    /// Returns the dual graph of the [`Triangulation`]
    pub fn dual_graph(&self) -> Graph {
        Graph::from_neighbour_iter(
            self.simplices
                .iter()
                .map(|simplex| simplex.get_neighbours().into_iter()),
        )
    }

    /// Returns an adjaceny list as nested vectors
    fn vertex_neighbours(&self) -> Vec<Vec<usize>> {
        let mut vertex_neighbours = vec![Vec::new(); self.num_vertices()];
        for (&va, &vb) in self
            .simplices
            .iter()
            .flat_map(|simplex| simplex.vertices.iter().tuple_combinations())
        {
            if va == vb {
                warn!(
                    "One-loop found in `Triangulation`. Function will likely not\
                        create correct `Graph` structure"
                );
                // If there is a loop the edge should appear twice in the graph, as
                // we are keeping track of outgoing edges
                vertex_neighbours[va].push(vb);
                vertex_neighbours[va].push(vb);
                continue;
            }
            if !vertex_neighbours[va].contains(&vb) {
                vertex_neighbours[va].push(vb);
            }
            if !vertex_neighbours[vb].contains(&va) {
                vertex_neighbours[vb].push(va);
            }
        }

        vertex_neighbours
    }

    /// Returns vertex graph of the [`Triangulation`]
    ///
    /// # Warning
    /// This function is only guaranteed to produce a correct vertex graph if the triangulation
    /// is a proper simplicial complex, i.e. non-degenerate where all simplices are given by
    /// unique sets of vertices.
    /// If this is not the case this function will still produce a graph but can have the
    /// incorrect multiplicity of links.
    pub fn vertex_graph(&self) -> Graph {
        Graph::from_neighbour_iter(
            self.vertex_neighbours()
                .into_iter()
                .map(|nbrs| nbrs.into_iter()),
        )
    }

    /// Returns vertex graph of the [`Triangulation`] explicity making all edges single edges.
    /// See [`vertex_graph()`](Self::vertex_graph) for further information.
    pub fn vertex_graph_noduplicates(&self) -> Graph {
        Graph::from_neighbour_iter(
            self.vertex_neighbours()
                .into_iter()
                .map(|nbrs| nbrs.into_iter().unique()),
        )
    }
}

impl<const N: usize> Triangulation<N> {
    /// Count the number of vertices by looking at all the vertex labels on each simplex.
    ///
    /// Note: if the [`Triangulation`] is a correct structure, this should be equal to
    /// `self.num_vertices()` which is just a direct lookup. Whereas this method can be slow.
    pub fn count_vertices(&self) -> usize {
        self.simplices
            .iter()
            .flat_map(|simplex| simplex.vertices.into_iter())
            .unique()
            .count()
    }

    /// Count the number of (M-1)-subsimplices, assuming the triangulation is a proper
    /// simplicial complex, i.e. is uniquely given by the vertices.
    ///
    /// This works by looping over all simplices and finding all unique
    /// unordered (v0, v1, ...) vertex tuples. This means this method will only
    /// return the correct number of edges if the triangulation is a proper
    /// simplicial complex, i.e. all simplices are uniquely given by their vertices.
    pub fn count_subsimplices<const M: usize>(&self) -> usize {
        self.simplices
            .iter()
            .flat_map(|simplex| simplex.vertices.into_iter().array_combinations())
            .map(canonical_simplex::<_, M>)
            .unique()
            .count()
    }

    /// Count the number of edges/1-simplices, see [`count_subsimplices()`](Self::count_subsimplices).
    pub fn count_n1(&self) -> usize {
        self.count_subsimplices::<2>()
    }

    /// Count the number of triangles/2-simplices, see [`count_subsimplices()`](Self::count_subsimplices).
    pub fn count_n2(&self) -> usize {
        self.count_subsimplices::<3>()
    }

    /// Count the number of tetrahedrons/3-simplices, see [`count_subsimplices()`](Self::count_subsimplices).
    pub fn count_n3(&self) -> usize {
        self.count_subsimplices::<4>()
    }

    /// Count the number of 4-simplices, see [`count_subsimplices()`](Self::count_subsimplices).
    pub fn count_n4(&self) -> usize {
        self.count_subsimplices::<5>()
    }
}

impl Triangulation2D {
    /// Computes the Euler characteristic of the triangulation
    ///
    /// This is given by `χ = N0 - N1 + N2` and is fixed for a given topology.
    /// For typical topolgies we have: χ(S2) = 2 or χ(T2) = 0
    pub fn euler_characteristic(&self) -> i32 {
        let n0 = self.num_vertices() as i32;
        let n1 = self.num_faces() as i32;
        let n2 = self.num_simplices() as i32;

        n0 - n1 + n2
    }
}

impl Triangulation3D {
    /// Computes the Euler characteristic of the triangulation
    ///
    /// This is given by `χ = N0 - N1 + N2 - N3` and is fixed for a given topology.
    /// For typical topolgies we have: χ(S3) = 2 or χ(T3) = 0
    pub fn euler_characteristic(&self) -> i32 {
        let n0 = self.num_vertices() as i32;
        let n1 = self.count_n1() as i32;
        let n2 = self.num_faces() as i32;
        let n3 = self.num_simplices() as i32;

        n0 - n1 + n2 - n3
    }
}

impl Triangulation4D {
    /// Computes the Euler characteristic of the triangulation
    ///
    /// This is given by `χ = N0 - N1 + N2 - N3 + N4` and is fixed for a given topology.
    /// For typical topolgies we have: χ(S4) = 2 or χ(T4) = 0
    pub fn euler_characteristic(&self) -> i32 {
        let n0 = self.num_vertices() as i32;
        let n1 = self.count_n1() as i32;
        let n2 = self.count_n2() as i32;
        let n3 = self.num_faces() as i32;
        let n4 = self.num_simplices() as i32;

        n0 - n1 + n2 - n3 + n4
    }
}

impl<const N: usize> Triangulation<N> {
    /// Swap orientation of D-simplex (N = D+1)
    ///
    /// This amount to swapping the ordering of the 0th and 1st vertex of the simplex
    /// and adapting all other data accordingly.
    pub fn flip_orientation(&mut self, label: usize) {
        // In old neighbour 0 swap backindex to be 1
        let (nbr0, backidx0) = self.get_neighbour_backindex(label, 0);
        self.simplices[nbr0].neighbours[backidx0] = N * label + 1; // Change Ni+0 to Ni+1

        // In old neighbour 1 swap backindex to be 0
        let (nbr1, backidx1) = self.get_neighbour_backindex(label, 1);
        self.simplices[nbr1].neighbours[backidx1] = N * label; // Change Ni+1 to Ni+0

        // Swap vertex 0 and vertex 1
        self.simplices[label].vertices.swap(0, 1);
        // Swap neighbour 0 and 1
        self.simplices[label].neighbours.swap(0, 1);
    }
}

impl<const N: usize> Triangulation<N>
where
    [usize; N]: FaceSimplex,
    Face<[usize; N]>: CanonicalSimplex + Hash + Eq + Debug,
{
    /// Reorient 4-simplices such that all simplices have an orientation consistent with
    /// the orientation of the the `origin` simplex.
    ///
    /// Think: generalization of the normal vector is consistent between all simplices.
    ///
    /// Note: makes use of vertex ordering to find orientation, so this cannot reorient properly
    /// when simplices have the same vertex more than once.
    pub fn reorient_simplices(&mut self, origin: usize) {
        // Perform a BFS
        let mut visited = vec![false; self.num_simplices()];
        visited[origin] = true;
        let mut queue = VecDeque::with_capacity(self.num_simplices() / 4);
        queue.push_back(origin);

        loop {
            let Some(node) = queue.pop_front() else {
                break;
            };
            for (idx, (nbr, backidx)) in self
                .get_neighbours_backindices(node)
                .into_iter()
                .enumerate()
            {
                if visited[nbr] {
                    continue;
                }
                visited[nbr] = true;
                queue.push_back(nbr);

                // Check if common face of `nbr` with `node`, and check orientation,
                // if the orientation is inconsistent flip the `nbr`.
                let vs = self.get_simplex_vertices(node);
                let (mut parity, face) = get_face(vs, idx).make_canonical_oriented();
                parity = !(parity ^ idx.is_multiple_of(2)); // Work in sign of face boundary
                let vs_nbr = self.get_simplex_vertices(nbr);
                let (mut parity_nbr, face_nbr) =
                    get_face(vs_nbr, backidx).make_canonical_oriented();
                parity_nbr = !(parity_nbr ^ backidx.is_multiple_of(2)); // Work in sign of face boundary
                assert_eq!(
                    face, face_nbr,
                    "Triangulation has incorrect vertex connectivity"
                );
                if parity == parity_nbr {
                    self.flip_orientation(nbr);
                }
            }
        }
    }
}

impl<const N: usize> Triangulation<N>
where
    [usize; N]: FaceSimplex,
    Face<[usize; N]>: CanonicalSimplex + Hash + Eq,
{
    /// Constructs a D-dimensional [`Triangulation<D+1>`] from a simplicial complex.
    ///
    /// The simplicial `complex` should be a set of ordered subsets of the vertices. In this case
    /// it means that every 4-simplex is given by 5 (unique!) vertices.
    /// Note: `complex` must encode a valid 4D triangulation, meaning that every 4-simplex is connected
    /// to exactly 5 other 4-simplices through its five faces (codim 1). Each tetrahedron can only
    /// have two 4-simplices connected to it. The current implementation does not support boundaries.
    /// Additionally this function requires vertex labels are dense, i.e. all in (0..num_vertices)
    /// If the `complex` gives an invalid structure this function may return an invalid structure.
    ///
    /// The order of the simplices in the iterator is used as the labels for the simplices
    pub fn from_simplicial_complex_iter<I: Iterator<Item = [usize; N]> + ExactSizeIterator>(
        complex: I,
    ) -> Self {
        // Maps face to (simplex label, face index)
        let mut face_to_simplex: HashMap<Face<[usize; N]>, (usize, usize)> =
            HashMap::with_capacity(complex.len());

        // New simplices
        let mut simplices = vec![
            Simplex {
                neighbours: [usize::MAX; N],
                vertices: [usize::MAX; N]
            };
            complex.len()
        ];
        for (slabel, simplex) in complex.enumerate() {
            // Set correct vertices for each simplex
            simplices[slabel].vertices = simplex;

            // Loop over faces with `face_index` that are opposite the corresponding vertex
            for (face_index, face) in get_faces(simplex)
                .into_iter()
                .map(|face| face.make_canonical())
                .enumerate()
            {
                // Check if face was already occured (should only happen once per face, hence del)
                if let Some((neighbour, back_index)) = face_to_simplex.remove(&face) {
                    // Set neighbour connections for self and the neighbour
                    simplices[slabel].neighbours[face_index] = N * neighbour + back_index;
                    simplices[neighbour].neighbours[back_index] = N * slabel + face_index;
                } else {
                    // For each edge insert first triangle occurence
                    face_to_simplex.insert(face, (slabel, face_index));
                }
            }
        }

        let mut max_vlabel = 0;
        for triangle in simplices.iter() {
            for vlabel in triangle.vertices {
                assert_ne!(vlabel, usize::MAX, "Some simplices have not been set");
                if vlabel > max_vlabel {
                    max_vlabel = vlabel;
                }
            }
            assert!(
                !triangle.neighbours.into_iter().any(|nbr| nbr == usize::MAX),
                "Not all labels have been set. The complex structure does encode a triangulation"
            );
        }

        Self {
            num_vertices: max_vlabel + 1,
            simplices,
        }
    }

    /// Constructs a [`Triangulation4D`] from a simplicial complex.
    ///
    /// The simplicial `complex` should be a set of ordered subsets of the vertices. In this case
    /// it means that every 4-simplex is given by 5 (unique!) vertices.
    /// Note: `complex` must encode a valid 4D triangulation, meaning that every 4-simplex is connected
    /// to exactly 5 other 4-simplices through its five faces (codim 1). Each tetrahedron can only
    /// have two 4-simplices connected to it. The current implementation does not support boundaries.
    /// Additionally this function requires vertex labels are dense, i.e. all in (0..num_vertices)
    /// If the `complex` gives an invalid structure this function may return an invalid structure.
    ///
    /// The order of the simplices in the slice is used as the labels for the simplices
    pub fn from_simplicial_complex(complex: &[[usize; N]]) -> Self {
        Self::from_simplicial_complex_iter(complex.iter().copied())
    }
}

impl Triangulation2D {
    /// Coordination numbers of number of 2-simplices around 0-simplices,
    /// i.e. around vertices and edges.
    ///
    /// Each coordination number list give the coordination number of the corresponding face.
    pub fn simplex_coordination(&self) -> Vec<usize> {
        let mut vertex_coordination: Vec<usize> = vec![0; self.num_vertices()];
        for simplex in self.simplices.iter() {
            for vertex in simplex.vertices.iter().copied() {
                vertex_coordination[vertex] += 1;
            }
        }

        vertex_coordination
    }
}

impl Triangulation3D {
    /// Coordination numbers of number of 3-simplices around (0, 1)-simplices,
    /// i.e. around vertices and edges.
    ///
    /// Each coordination number list give the coordination number of the corresponding face.
    pub fn simplex_coordination(&self) -> (Vec<usize>, Vec<usize>) {
        let mut vertex_coordination: Vec<usize> = vec![0; self.num_vertices()];
        let mut edge_coordination: HashMap<[usize; 2], usize> = HashMap::new();
        for simplex in self.simplices.iter() {
            for vertex in simplex.vertices.iter().copied() {
                vertex_coordination[vertex] += 1;
            }
            for edge in simplex.vertices.iter().copied().array_combinations() {
                *edge_coordination.entry(edge).or_default() += 1;
            }
        }
        let edge_coordination = edge_coordination.into_values().collect();

        (vertex_coordination, edge_coordination)
    }
}

impl Triangulation4D {
    /// Coordination numbers of number of 4-simplices around (0, 1, 2)-simplices,
    /// i.e. around vertices, edges, and triangles.
    ///
    /// Each coordination number list give the coordination number of the corresponding face.
    pub fn simplex_coordination(&self) -> (Vec<usize>, Vec<usize>, Vec<usize>) {
        let mut vertex_coordination: Vec<usize> = vec![0; self.num_vertices()];
        let mut edge_coordination: HashMap<[usize; 2], usize> = HashMap::new();
        let mut triangle_coordination: HashMap<[usize; 3], usize> = HashMap::new();
        for simplex in self.simplices.iter() {
            for vertex in simplex.vertices.iter().copied() {
                vertex_coordination[vertex] += 1;
            }
            for edge in simplex.vertices.iter().copied().array_combinations() {
                *edge_coordination.entry(edge).or_default() += 1;
            }
            for triangle in simplex.vertices.iter().copied().array_combinations() {
                *triangle_coordination.entry(triangle).or_default() += 1;
            }
        }
        let edge_coordination = edge_coordination.into_values().collect();
        let triangle_coordination = triangle_coordination.into_values().collect();

        (
            vertex_coordination,
            edge_coordination,
            triangle_coordination,
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use test_log::test;

    #[test]
    fn test_from_simpcomp_tetra() {
        // A simplicial complex of a tetrahedron
        const TETRA_COMPLEX_STR: &str = "021 123 013 203";
        let complex = TETRA_COMPLEX_STR
            .split(' ')
            .map(|triplet| {
                triplet
                    .chars()
                    .map(|c| c.to_digit(10).unwrap() as usize)
                    .collect_array()
                    .unwrap()
            })
            .collect_vec();

        println!("{complex:?}");
        let triangulation = Triangulation2D::from_simplicial_complex(&complex);
        println!("{triangulation:?}");

        triangulation.check_simplex_counts().unwrap();
        triangulation.check_vertex_consistency().unwrap();
        triangulation.check_vertex_consistency_ordered().unwrap();
        triangulation.check_biconnectivity().unwrap();

        let n0 = triangulation.num_vertices();
        let n1 = triangulation.num_faces();
        let n2 = triangulation.num_simplices();
        println!("N0: {}, N1: {}, N2: {}", n0, n1, n2);

        let chi = triangulation.euler_characteristic();
        assert_eq!(chi, 2, "Euler charecteristic is incorrect");
    }

    #[test]
    fn test_from_simpcomp_torus() {
        // A simplicial complex of a torus
        const TORUS_COMPLEX_STR: &str =
            "014 043 125 154 203 235 347 376 458 487 536 568 671 610 782 721 860 802";
        let complex = TORUS_COMPLEX_STR
            .split(' ')
            .map(|triplet| {
                triplet
                    .chars()
                    .map(|c| c.to_digit(10).unwrap() as usize)
                    .collect_array()
                    .unwrap()
            })
            .collect_vec();

        let triangulation = Triangulation2D::from_simplicial_complex(&complex);

        triangulation.check_simplex_counts().unwrap();
        triangulation.check_vertex_consistency().unwrap();
        triangulation.check_vertex_consistency_ordered().unwrap();
        triangulation.check_biconnectivity().unwrap();

        let n0 = triangulation.num_vertices();
        let n1 = triangulation.num_faces();
        let n2 = triangulation.num_simplices();
        println!("N0: {}, N1: {}, N2: {}", n0, n1, n2);

        let chi = triangulation.euler_characteristic();
        assert_eq!(chi, 0, "Euler charecteristic is incorrect");
    }

    #[test]
    fn test_reorientation() {
        const TORUS_COMPLEX_STR: &str =
            "014 043 125 154 203 235 347 376 458 487 536 568 671 610 782 721 860 802";
        let complex = TORUS_COMPLEX_STR
            .split(' ')
            .map(|triplet| {
                triplet
                    .chars()
                    .map(|c| c.to_digit(10).unwrap() as usize)
                    .collect_array()
                    .unwrap()
            })
            .collect_vec();

        let mut triangulation = Triangulation2D::from_simplicial_complex(&complex);

        triangulation.check_vertex_consistency_ordered().unwrap();

        triangulation.flip_orientation(3);
        triangulation
            .check_vertex_consistency_ordered()
            .unwrap_err();

        triangulation.reorient_simplices(3);
        triangulation.check_vertex_consistency_ordered().unwrap();
    }
}