dyntri-core 0.10.2

use std::collections::HashMap;
use std::hash::Hash;

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

use crate::graph::Graph;
use crate::triangulation::simplices::{
    CausalSimplex, CausalSimplex2Kind, CausalSimplex3Kind, CausalSimplex4Kind, CausalSimplexKind,
    canonical_simplex,
};
use crate::triangulation::triangulation::Triangulation;

/// Triangulation structure designed to store a causal "gluing" of (N-1)-simplices.
///
/// Internally four half-faces codim 1 (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 CausalTriangulation<K, const N: usize> {
    /// Number of time-slices / spatial slices in the triangulation
    pub ntime: u16,
    /// The time labels of each vertex
    ///
    /// These time labels must be such that the lower vertices of a simplex have the same
    /// time label as the simplex itself
    pub vertex_times: Vec<u16>,
    /// List holding (N-1)-simplices of the triangulation.
    ///
    /// Note that the `neighbours` field of each (N-2)-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`
    ///
    /// If the triangulation if absolutely ordered (i.e. [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`),
    /// this means the timelike neighbour (if it exists) is always the 0th elements in the list.
    pub simplices: Vec<CausalSimplex<K, N>>,
}

pub type CausalTriangulation2D = CausalTriangulation<CausalSimplex2Kind, 3>;
pub type CausalTriangulation3D = CausalTriangulation<CausalSimplex3Kind, 4>;
pub type CausalTriangulation4D = CausalTriangulation<CausalSimplex4Kind, 5>;

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

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Return the number of timeslices in the triangulation
    #[inline]
    pub fn num_timeslices(&self) -> u16 {
        self.ntime
    }

    /// Returns the number of vertices in the triangulation
    #[inline]
    pub fn num_vertices(&self) -> usize {
        self.vertex_times.len()
    }

    /// 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`](super::simplices::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<K: Copy, const N: usize> CausalTriangulation<K, N> {
    /// Convert to non-causal/Euclidean triangulation by copying from references
    pub fn to_triangulation(&self) -> Triangulation<N> {
        Triangulation {
            num_vertices: self.num_vertices(),
            simplices: self.simplices.iter().copied().map(From::from).collect(),
        }
    }
}

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Convert to non-causal/Euclidean triangulation
    pub fn into_triangulation(self) -> Triangulation<N> {
        Triangulation {
            num_vertices: self.num_vertices(),
            simplices: self.simplices.into_iter().map(From::from).collect(),
        }
    }
}

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Returns a reference to time labels of the vertices of the [`CausalTriangulation`]
    #[inline]
    pub fn get_vertex_times(&self) -> &[u16] {
        &self.vertex_times
    }

    /// Return the time label of the vertex given by `label`
    #[inline]
    pub fn get_vertex_time(&self, label: usize) -> u16 {
        self.vertex_times[label]
    }

    /// Return the time label of the (N-1)-simplex given by `label`
    #[inline]
    pub fn get_simplex_time(&self, label: usize) -> u16 {
        self.get_simplex(label).t
    }

    /// Returns a reference to the (N-1)-simplex list of the [`CausalTriangulation`]
    #[inline]
    pub fn get_simplices(&self) -> &[CausalSimplex<K, 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 [`CausalTriangulation`].
    ///
    /// # Panic
    /// This method will panic if `label >= [self.num_simplices()]`
    #[inline]
    pub fn get_simplex(&self, label: usize) -> &CausalSimplex<K, N> {
        &self.simplices[label]
    }

    /// Returns a mutable reference to the (N-1)-simplex given by `label`
    ///
    /// See [`CausalTriangulation::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 CausalSimplex<K, 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<K: Copy, const N: usize> CausalTriangulation<K, N> {
    /// Return the kind of the (N-1)-simplex given by `label`
    ///
    /// See e.g. [CausalSimplex4Kind](crate::triangulation::simplices::CausalSimplex4Kind)
    #[inline]
    pub fn get_simplex_kind(&self, label: usize) -> K {
        self.get_simplex(label).kind
    }
}

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Returns the time-slice/spatial slab volume profile,
    /// i.e. the number of vertices in each timeslice.
    pub fn timeslice_volume_profile(&self) -> Vec<usize> {
        let mut profile = vec![0; self.ntime as usize];
        for t in self.vertex_times.iter().copied() {
            profile[t as usize] += 1;
        }
        profile
    }

    /// Returns the spatial slab volume profile, i.e. the number of simplices in each
    /// slab/thick timeslice.
    ///
    /// Note: for 2D CDT this should be equal to the moving sum of the
    /// [`timeslice_volume_profile()`](Self::timeslice_volume_profile) with a window of 2.
    pub fn timeslab_volume_profile(&self) -> Vec<usize> {
        let mut profile = vec![0; self.ntime as usize];
        for triangle in &self.simplices {
            profile[triangle.t as usize] += 1;
        }
        profile
    }
}

impl CausalTriangulation2D {
    /// Returns the fractional spatial volume profile
    ///
    /// Fractional here means the consecutive volume in each slab, e.g.
    /// 2,1-type; and 1,2-type; 2-simplices.
    ///
    /// Note the number of 2,1-type and 1,2-type triangles is equal to eachother
    /// in adjacent layers in a correct triangulation.
    pub fn timeslab_fractional_volume_profile(&self) -> Vec<usize> {
        let mut profile = vec![0; 2 * (self.ntime as usize)];
        for simplex in &self.simplices {
            type SimplexKind = super::simplices::CausalSimplex2Kind;
            match simplex.kind {
                SimplexKind::S21 => profile[simplex.t as usize] += 1,
                SimplexKind::S12 => profile[simplex.t as usize + 1] += 1,
            }
        }
        profile
    }
}

impl CausalTriangulation3D {
    /// Returns the fractional spatial volume profile
    ///
    /// Fractional here means the consecutive volume in each slab, e.g.
    /// 3,1-type; 2,2-type; and 1,3-type; 3-simplices.
    ///
    /// Note the number of 3,1-type and 1,3-type triangles is equal to eachother
    /// in adjacent layers in a correct triangulation.
    pub fn timeslab_fractional_volume_profile(&self) -> Vec<usize> {
        let mut profile = vec![0; 3 * (self.ntime as usize)];
        for simplex in &self.simplices {
            type SimplexKind = super::simplices::CausalSimplex3Kind;
            match simplex.kind {
                SimplexKind::S31 => profile[simplex.t as usize] += 1,
                SimplexKind::S22 => profile[simplex.t as usize + 1] += 1,
                SimplexKind::S13 => profile[simplex.t as usize + 2] += 1,
            }
        }
        profile
    }
}

impl CausalTriangulation4D {
    /// Returns the fractional spatial volume profile
    ///
    /// Fractional here means the consecutive volume in each slab, e.g.
    /// 4,1-type; 3,2-type; 2,3-type; and 1,4-type 4-simplices.
    ///
    /// Note the number of 4,1-type and 1,4-type triangles is equal to eachother
    /// in adjacent layers in a correct triangulation.
    pub fn timeslab_fractional_volume_profile(&self) -> Vec<usize> {
        let mut profile = vec![0; 4 * (self.ntime as usize)];
        for simplex in &self.simplices {
            type SimplexKind = super::simplices::CausalSimplex4Kind;
            match simplex.kind {
                SimplexKind::S41 => profile[simplex.t as usize] += 1,
                SimplexKind::S32 => profile[simplex.t as usize + 1] += 1,
                SimplexKind::S23 => profile[simplex.t as usize + 2] += 1,
                SimplexKind::S14 => profile[simplex.t as usize + 3] += 1,
            }
        }
        profile
    }
}

impl<K, const N: usize> CausalTriangulation<K, N> {
    /// Returns oriented dual graph of the [`CausalTriangulation`]
    ///
    /// Note that this graph is oriented in the sense that the links are given in the same
    /// order as the neighbour relation were orderedin the [CausalTriangulation].
    /// In 2D this ought to be in counter-clockwise(positive) order, when used as designed.
    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 [`CausalTriangulation`]
    ///
    /// # 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 [`CausalTriangulation`] 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<K, const N: usize> CausalTriangulation<K, N> {
    /// Count the number of vertices by looking at all the vertex labels on each simplex.
    ///
    /// Note: if the [`CausalTriangulation`] 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<K: Hash + Ord + Copy, const N: usize> CausalTriangulation<K, N> {
    /// Returns the counts of (N-1)-simplices as a map from the SimplexKind to count
    ///
    /// This works by looping over all simplices O(N)
    pub fn count_simplex_kinds(&self) -> HashMap<K, usize> {
        let mut counts = HashMap::with_capacity(N - 1);
        for simplex in &self.simplices {
            *counts.entry(simplex.kind).or_insert(0) += 1;
        }

        counts
    }
}

impl CausalTriangulation2D {
    /// 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 CausalTriangulation3D {
    /// 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 CausalTriangulation4D {
    /// 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<K: Copy + CausalSimplexKind, const N: usize> CausalTriangulation<K, N> {
    /// Gets the neighbour in the future direction of the simplex `label`.
    ///
    /// This returns the neighbour label if the simplex has a future neighbour,
    /// or if depending on its type it has no future neighbour, it returns `None`
    ///
    /// Note: This only gives the correct result if
    /// [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`.
    pub fn get_future_neighbour(&self, label: usize) -> Option<usize> {
        match self.get_simplex_kind(label).orientation() {
            super::simplices::CausalOrientation::Past => None,
            super::simplices::CausalOrientation::Space => None,
            super::simplices::CausalOrientation::Future => Some(self.get_neighbour(label, 0)),
        }
    }

    /// Gets the neighbour in the past direction of the simplex `label`.
    ///
    /// This returns the neighbour label if the simplex has a past neighbour,
    /// or if depending on its type it has no past neighbour, it returns `None`
    ///
    /// Note: This only gives the correct result if
    /// [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`.
    pub fn get_past_neighbour(&self, label: usize) -> Option<usize> {
        match self.get_simplex_kind(label).orientation() {
            super::simplices::CausalOrientation::Past => Some(self.get_neighbour(label, 0)),
            super::simplices::CausalOrientation::Space => None,
            super::simplices::CausalOrientation::Future => None,
        }
    }

    /// Gets the neighbour in the timelike direction of the simplex `label`.
    ///
    /// This returns the neighbour label if the simplex has a timelike neighbour,
    /// or if depending on its type it has no timelike neighbour, it returns `None`.
    /// In 2D there are only 2,1 and 1,2 type triangles, which always have a timelike neighbour.
    ///
    /// Note: This only gives the correct result if
    /// [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`.
    pub fn get_timelike_neighbour(&self, label: usize) -> Option<usize> {
        match self.get_simplex_kind(label).orientation() {
            super::simplices::CausalOrientation::Past => Some(self.get_neighbour(label, 0)),
            super::simplices::CausalOrientation::Space => None,
            super::simplices::CausalOrientation::Future => Some(self.get_neighbour(label, 0)),
        }
    }

    /// Returns iterator over the spacelike neighbours of the simplex `label`.
    ///
    /// This returns only the neighbour labels in the spacelike direction, the number of which
    /// depends on the type of the simplex.
    /// In 2D there are only 2,1 and 1,2 type triangles, which always have 2 spacelike neighbours.
    ///
    /// Note: This only gives the correct result if
    /// [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`.
    pub fn get_spacelike_neighbours(&self, label: usize) -> impl Iterator<Item = usize> {
        let simplex = self.get_simplex(label);
        type Orientation = super::simplices::CausalOrientation;
        match simplex.kind.orientation() {
            Orientation::Past | Orientation::Future => simplex.get_neighbours().into_iter().skip(1),
            #[allow(clippy::iter_skip_zero)]
            Orientation::Space => simplex.get_neighbours().into_iter().skip(0),
        }
    }
}

impl CausalTriangulation2D {
    /// Returns the left spacelike neighbour of the simplex `label`, w.r.t future being up.
    ///
    /// Note: This only gives the correct result if
    /// [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`.
    pub fn get_left_neighbour(&self, label: usize) -> usize {
        let triangle = self.get_simplex(label);
        match triangle.kind {
            super::simplices::CausalSimplex2Kind::S21 => triangle.get_neighbour(2),
            super::simplices::CausalSimplex2Kind::S12 => triangle.get_neighbour(1),
        }
    }

    /// Returns the right spacelike neighbour of the simplex `label`, w.r.t future being up.
    ///
    /// Note: This only gives the correct result if
    /// [`CausalTriangulation::check_neighbour_absolute_ordering`] is `Ok`.
    pub fn get_right_neighbour(&self, label: usize) -> usize {
        let triangle = self.get_simplex(label);
        match triangle.kind {
            super::simplices::CausalSimplex2Kind::S21 => triangle.get_neighbour(1),
            super::simplices::CausalSimplex2Kind::S12 => triangle.get_neighbour(2),
        }
    }
}