dyntri-core 0.10.2

use std::collections::{HashMap, VecDeque};

use itertools::{Itertools, izip};
use log::trace;

use crate::graph::Graph;
use crate::triangulation::simplices::{CausalSimplex3Kind, CausalSimplex4Kind, Simplex2, Simplex3};
use crate::triangulation::timeslices::sparsemap::SparseToDenseMap;
use crate::triangulation::{
    CausalTriangulation, CausalTriangulation2D, CausalTriangulation3D, CausalTriangulation4D,
    Triangulation2D, Triangulation3D,
};

impl CausalTriangulation2D {
    /// Walk along triangles until the original triangle is reached
    ///
    /// Restrict to stay between the origin timeslice and `ntime` further timeslices, find
    /// the shortest loop under this restriction.
    ///
    /// Note: the returned path is in the past direction, ending not stating with the `origin`
    pub fn identify_timelike_loop(&self, origin: usize) -> Vec<usize> {
        let origin_triangle = self.get_simplex(origin);
        // Make sure the starting triangle is of 21-type, get neighbour if not.
        type TriangleKind = super::simplices::CausalSimplex2Kind;
        let (origin, origin_triangle) = match origin_triangle.kind {
            TriangleKind::S21 => (origin, origin_triangle),
            TriangleKind::S12 => {
                let upper = origin_triangle.get_neighbour(0);
                (upper, self.get_simplex(upper))
            }
        };
        let t_origin = origin_triangle.t;

        let mut visited = vec![None; self.num_simplices()];
        visited[origin] = Some(usize::MAX);

        let mut queue = VecDeque::with_capacity(2 * self.num_simplices() / (self.ntime as usize));
        queue.push_back(origin);

        'outer: loop {
            let label = queue
                .pop_front()
                .expect("Loop should always be found in proper CDT");
            if let Some(nbr) = self.get_future_neighbour(label) {
                if nbr == origin {
                    // Went back to original triangle
                    visited[nbr] = Some(label);
                    break 'outer;
                }

                if self.get_simplex_time(nbr) != t_origin && visited[nbr].is_none() {
                    visited[nbr] = Some(label);
                    {
                        queue.push_back(nbr);
                    }
                }
            }
            for nbr in self.get_spacelike_neighbours(label) {
                // Only consider unvisited nodes
                if visited[nbr].is_some() {
                    continue;
                }
                visited[nbr] = Some(label);
                queue.push_back(nbr);
            }
        }

        // Rewalk path
        let mut path = Vec::with_capacity((self.ntime * 2).into());
        let mut current_triangle = origin;
        loop {
            current_triangle = visited[current_triangle]
                .expect("The path should have been fully explored by construction");
            path.push(current_triangle);
            if current_triangle == origin {
                break;
            }
        }

        path
    }
}

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Returns a nested list of all simplex labels in each time-slab
    ///
    /// The outer list is over time-slabs, the inner list over simplices in arbitrary order.
    pub fn timeslab_simplices(&self) -> Vec<Vec<usize>> {
        let mean_timeslab_len = self.num_simplices() / self.num_timeslices() as usize;
        let mut timeslice_labels: Vec<_> =
            vec![Vec::with_capacity(mean_timeslab_len); self.num_timeslices() as usize];
        for (label, simplex) in self.get_simplices().iter().enumerate() {
            let t = simplex.t as usize;
            timeslice_labels[t].push(label);
        }
        timeslice_labels
    }

    /// Returns a nested list of all vertex labels in each time-slice
    ///
    /// The outer list is over time-slices, the inner list over vertices in arbitrary order.
    pub fn timeslice_vertices(&self) -> Vec<Vec<usize>> {
        let mean_timeslice_len = self.num_vertices() / self.num_timeslices() as usize;
        let mut timeslice_labels: Vec<_> =
            vec![Vec::with_capacity(mean_timeslice_len); self.num_timeslices() as usize];
        for (label, &t) in self.vertex_times.iter().enumerate() {
            timeslice_labels[t as usize].push(label);
        }
        timeslice_labels
    }
}

impl CausalTriangulation4D {
    /// Construct the [`Triangulation3D`] that is the time-slice of [`CausalTriangulation4D`] at `t`
    pub fn timeslice_triangulation(&self, t: u16) -> Triangulation3D {
        if t >= self.num_timeslices() {
            panic!("t={t} is out of bound of the triangulation");
        }
        // Get a map for the vertex labels to new 'dense' labels
        let mut remap = SparseToDenseMap::new(self.num_vertices());

        trace!("Get tetrahedron vertices");
        let tetras: Vec<[usize; 4]> = self
            .simplices
            .iter()
            .filter(|simplex| simplex.t == t)
            .filter(|simplex| matches!(simplex.kind, CausalSimplex4Kind::S41))
            .map(|simplex| {
                let terta_vertices: [usize; 4] = simplex
                    .vertices
                    .iter()
                    .copied()
                    .filter(|&vlabel| self.get_vertex_time(vlabel) == simplex.t)
                    .map(|vlabel| remap.add(vlabel))
                    .collect_array()
                    .expect("There must be only 4 vertices with correct time, as kind is S41");
                terta_vertices
            })
            .collect();

        trace!("Reconstructing tetrahedron connectivity");
        let mut triangles: HashMap<[usize; 3], usize> = HashMap::with_capacity(tetras.len() * 2);
        let mut neighbours: Vec<Vec<usize>> = vec![Vec::with_capacity(4); tetras.len()];
        for (slabel, tetra) in tetras.iter().enumerate() {
            for mut triangle in tetra.iter().copied().array_combinations() {
                // Sort the triangle vertices to get unique identifier
                triangle.sort_unstable();
                // Is triangle was already present we can make the links of the tertaeder
                if let Some(slabel_nbr) = triangles.insert(triangle, slabel) {
                    neighbours[slabel_nbr].push(slabel);
                    neighbours[slabel].push(slabel_nbr);
                };
            }
        }

        trace!("Construct 3-simplices");
        assert_eq!(tetras.len(), neighbours.len());
        let simplices: Vec<Simplex3> = izip!(tetras, neighbours)
            .map(|(tetra, neighbour)| Simplex3 {
                vertices: tetra,
                neighbours: neighbour
                    .try_into()
                    .expect("We should find 4 neighbours of each tetraeder"),
            })
            .collect();

        Triangulation3D {
            num_vertices: remap.len(),
            simplices,
        }
    }
}

impl CausalTriangulation3D {
    /// Construct the [`Triangulation2D`] that is the time-slice of [`CausalTriangulation3D`] at `t`
    ///
    /// Note: the orientation of the simplices of such a timeslices is not guaranteed to be proper consistent,
    /// even if the original `CausalTriangulation` is. If proper orientation is desired use:
    /// ```triangulation.timeslices_trianglation(t).reorient_simplices(any_vertex);```
    pub fn timeslice_triangulation(&self, t: u16) -> Triangulation2D {
        if t >= self.num_timeslices() {
            panic!("t={t} is out of bound of the triangulation");
        }
        // Get a map for the vertex labels to new 'dense' labels
        let mut remap = SparseToDenseMap::new(self.num_vertices());

        trace!("Get tetrahedron vertices");
        let triangles: Vec<_> = self
            .simplices
            .iter()
            .filter(|simplex| simplex.t == t)
            .filter(|simplex| matches!(simplex.kind, CausalSimplex3Kind::S31))
            .map(|simplex| {
                let triangle_vertices: [usize; 3] = simplex
                    .vertices
                    .iter()
                    .copied()
                    .filter(|&vlabel| self.get_vertex_time(vlabel) == simplex.t)
                    .map(|vlabel| remap.add(vlabel))
                    .collect_array()
                    .expect("There must be only 3 vertices with correct time, as kind is S31");
                triangle_vertices
            })
            .collect();

        trace!("Reconstructing triangle connectivity");
        let mut edges: HashMap<[usize; 2], usize> = HashMap::with_capacity(triangles.len() * 2);
        let mut neighbours: Vec<Vec<usize>> = vec![Vec::with_capacity(3); triangles.len()];
        for (slabel, triangle) in triangles.iter().enumerate() {
            for mut edge in triangle.iter().copied().array_combinations() {
                // Sort the edge vertices to get unique identifier
                edge.sort_unstable();
                // Is edge was already present we can make the connection between the triangles
                if let Some(slabel_nbr) = edges.insert(edge, slabel) {
                    neighbours[slabel_nbr].push(slabel);
                    neighbours[slabel].push(slabel_nbr);
                };
            }
        }

        trace!("Construct 2-simplices");
        let simplices: Vec<Simplex2> = izip!(triangles, neighbours)
            .map(|(triangle, neighbour)| Simplex2 {
                vertices: triangle,
                neighbours: neighbour
                    .try_into()
                    .expect("We should find 3 neighbours of each triangle"),
            })
            .collect();

        Triangulation2D {
            num_vertices: remap.len(),
            simplices,
        }
    }
}

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Construct the [`Graph`] that is the slab/thick time-slice of [`CausalTriangulation`]
    /// between `t` and `t+1`.
    pub fn construct_timeslab_dual(&self, t: u16) -> Graph {
        let mut label_map = SparseToDenseMap::new(self.num_simplices());

        let mean_timeslab_len = self.num_simplices() / self.num_timeslices() as usize;
        let mut nodes: Vec<Vec<_>> = Vec::with_capacity(mean_timeslab_len);

        for (label, simplex) in self.simplices.iter().enumerate() {
            if simplex.t != t {
                continue;
            }
            let new_label = label_map.add(label);
            let nbrs: Vec<_> = simplex
                .get_neighbours()
                .into_iter()
                .filter(|&nbr| self.get_simplex_time(nbr) == simplex.t)
                .map(|nbr| label_map.add(nbr))
                .collect();
            if new_label >= nodes.len() {
                nodes.resize_with(new_label + 1, Vec::new);
            }
            nodes[new_label] = nbrs;
        }

        Graph::from_neighbour_iter(nodes.into_iter().map(|inner| inner.into_iter()))
    }

    /// Construct the [`Graph`]s of the slabs/thick time-slices of [`CausalTriangulation`]
    /// between `t` and `t+1` for all `t`, starting with `t = 0`
    pub fn construct_all_timeslab_dual(&self) -> Vec<Graph> {
        let nsimplices = self.num_simplices();
        let ntime = self.num_timeslices();

        // Map to store original simplex label to dense new node labels (for each timeslab combined)
        let mut label_map = vec![usize::MAX; nsimplices];

        // Create mapping, storing back-labels in `timeslab_labels`
        let mean_timeslab_len = nsimplices / ntime as usize;
        let mut timeslab_labels: Vec<_> =
            vec![Vec::with_capacity(mean_timeslab_len); ntime as usize];
        for (label, simplex) in self.get_simplices().iter().enumerate() {
            let t = simplex.t as usize;
            label_map[label] = timeslab_labels[t].len();
            timeslab_labels[t].push(label);
        }

        let slab_to_graph = |t| {
            let adjacency = timeslab_labels[t as usize].iter().map(|&label_old| {
                self.get_neighbours(label_old)
                    .into_iter()
                    .filter(|&nbr| self.get_simplex_time(nbr) == t)
                    .map(|nbr| label_map[nbr])
            });
            Graph::from_neighbour_iter(adjacency)
        };

        (0..ntime).map(slab_to_graph).collect()
    }
}

mod sparsemap {
    #[derive(Debug, Clone)]
    pub(super) struct SparseToDenseMap {
        sparse: Vec<Option<usize>>,
        next_dense: usize,
    }

    impl SparseToDenseMap {
        pub(super) fn new(max_sparse_index: usize) -> Self {
            SparseToDenseMap {
                sparse: vec![None; max_sparse_index + 1],
                next_dense: 0,
            }
        }

        pub(super) fn len(&self) -> usize {
            self.next_dense
        }

        /// Returns the dense index for `sparse_index` creating a new index if it had not been set yet.
        pub(super) fn add(&mut self, sparse_index: usize) -> usize {
            let value = self
                .sparse
                .get_mut(sparse_index)
                .expect("Given sparse_index is larger than max_sparse_index");
            *value.get_or_insert_with(|| {
                let dense_value = self.next_dense;
                self.next_dense += 1;
                dense_value
            })
        }
    }
}