dyntri-core 0.10.2

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

#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
pub struct Graph {
    node_index: Vec<usize>,
    neighbours: Vec<usize>,
}

impl Graph {
    #[inline]
    /// Alias of [`num_nodes()`](Graph::num_nodes())
    pub fn len(&self) -> usize {
        self.num_nodes()
    }

    #[inline]
    /// Number of nodes/vertices in the [`Graph`]
    pub fn num_nodes(&self) -> usize {
        self.node_index.len() - 1
    }

    #[inline]
    /// Number of links in the [`Graph`]
    pub fn num_links(&self) -> usize {
        self.neighbours.len() / 2
    }

    #[inline]
    /// Returns if the [`Graph`] has no nodes
    pub fn is_empty(&self) -> bool {
        self.node_index.is_empty()
    }

    #[inline]
    /// Returns a slice of the neighbours of the node given by `label` in the [`Graph`]
    pub fn get_neighbours(&self, label: usize) -> &[usize] {
        self.get_neighbours_checked(label)
            .expect("`label` is out of bounds of the Graph")
    }

    /// Returns a slice of the neighbours of the node given by `label` in the [`Graph`]
    #[inline]
    pub fn get_neighbours_checked(&self, label: usize) -> Option<&[usize]> {
        if label + 1 >= self.node_index.len() {
            return None;
        }
        let istart = unsafe { *self.node_index.get_unchecked(label) };
        let iend = unsafe { *self.node_index.get_unchecked(label + 1) };
        let nbrs = self
            .neighbours
            .get(istart..iend)
            .expect("Correct graph structure must have neighbours of every node.");
        Some(nbrs)
    }

    #[inline]
    /// Returns the degree, i.e. number of neighbours of the node given by `label` in the [`Graph`]
    pub fn get_degree(&self, label: usize) -> usize {
        self.get_degree_checked(label)
            .expect("`label` is out of bounds of the Graph")
    }

    #[inline]
    /// Returns the degree, i.e. number of neighbours of the node given by `label` in the [`Graph`]
    pub fn get_degree_checked(&self, label: usize) -> Option<usize> {
        if label + 1 >= self.node_index.len() {
            return None;
        }
        let degree = unsafe {
            self.node_index.get_unchecked(label + 1) - self.node_index.get_unchecked(label)
        };
        Some(degree)
    }
}

impl Graph {
    /// Create a new [`Graph`] from its raw compressed sparse parts directly
    ///
    /// `neighbours` is a list specifying the neighbours of all nodes.
    /// `node_index` specifies the indices of the nodes in the `neighbours` list, such that:
    /// `neighbours[node_index[i]..node_index[i + 1]]` gives the neighbours of the node `i`.
    /// This means they must have the following properties:
    /// 1. `node_index[0]` == 0
    /// 2. `node_index.len() == num_nodes + 1`
    /// 3. `node_index[num_nodes] == neighbours.len()`
    ///
    /// Additionally, the every label that appears in `neighbours` corresponds to the nodes, hence
    /// 4. `max(neighbours) < num_nodex`.
    ///
    /// Finally, some functions assume the graph is undirected, meaning that the neighbour of each
    /// node should also have that node as neighbour.
    ///
    /// # Safety
    /// This creation method does checks if conditions (1) and (3) are obeyed, but does not
    /// guarentee the full validity of the created [`Graph`].
    /// It is up to the user to create a valid structure, if failed other fuctions using the
    /// [`Graph`] could give incorrect results of panic.
    #[inline]
    pub fn from_raw_parts(node_index: Vec<usize>, neighbours: Vec<usize>) -> Option<Self> {
        if let Some(1..) = node_index.first() {
            return None;
        }

        match node_index.last() {
            Some(&idx) if idx != neighbours.len() => return None,
            _ => (),
        }

        Some(Self {
            node_index,
            neighbours,
        })
    }

    /// Create a new [`Graph`] from without performing any validity checks,
    /// see [`from_raw_parts()`](Self::from_raw_parts).
    #[inline]
    pub fn from_raw_parts_unchecked(node_index: Vec<usize>, neighbours: Vec<usize>) -> Self {
        Self {
            node_index,
            neighbours,
        }
    }

    /// Create a new [`Graph`] from given adjacency/neighbour relations.
    ///
    /// The `adjacency` is given through nested iterators, where the inner iterator
    /// iterates over all neighbour labels of each nodes, and the outer iterator iterates
    /// over all nodes in the graph.
    ///
    /// The labels are expected to be `usize` and must be such that the `k`th element in
    /// outer iterator refers to the node with label `k`.
    ///
    /// See also [`from_neighbour_slice()`](Self::from_neighbour_slice)
    pub fn from_neighbour_iter<I, J>(adjacency: I) -> Self
    where
        I: Iterator<Item = J>,
        J: Iterator<Item = usize>,
    {
        let mut node_index = Vec::with_capacity(adjacency.size_hint().0);
        let mut neighbours = Vec::with_capacity(adjacency.size_hint().0 * 6);
        for nbrs in adjacency {
            node_index.push(neighbours.len());
            neighbours.extend(nbrs);
        }
        node_index.push(neighbours.len());

        Self {
            node_index,
            neighbours,
        }
    }

    /// Create a new [`Graph`] from given adjacency/neighbour relations.
    ///
    /// The `adjacency` is given through an iterator over neighbour slices. The iterator
    /// iterates over all nodes in the graph and gives the neighbours of each of them as slice.
    ///
    /// The labels are expected to be `usize` and must be such that the `k`th element in
    /// outer iterator refers to the node with label `k`.
    ///
    /// See also [`from_neighbour_iter()`](Self::from_neighbour_iter)
    pub fn from_neighbour_slice<'a, I, S>(adjacency: I) -> Self
    where
        I: Iterator<Item = &'a S>,
        S: AsRef<[usize]> + ?Sized + 'a,
    {
        let mut node_index = Vec::with_capacity(adjacency.size_hint().0);
        let mut neighbours = Vec::with_capacity(adjacency.size_hint().0 * 6);
        for nbrs in adjacency {
            node_index.push(neighbours.len());
            neighbours.extend_from_slice(nbrs.as_ref());
        }
        node_index.push(neighbours.len());

        Self {
            node_index,
            neighbours,
        }
    }
}

impl Graph {
    #[inline]
    /// Returns a iterator over the slices of neighbours of each node
    pub fn iter(&'_ self) -> GraphIter<'_> {
        GraphIter::new(self)
    }
}

#[derive(Debug, Clone, Copy)]
pub struct GraphIter<'a> {
    graph: &'a Graph,
    label: usize,
    istart: usize,
    iend: usize,
}

impl<'a> GraphIter<'a> {
    fn new(graph: &'a Graph) -> Self {
        Self {
            graph,
            label: 0,
            istart: 0,
            iend: 0,
        }
    }
}

impl<'a> Iterator for GraphIter<'a> {
    type Item = &'a [usize];

    fn next(&mut self) -> Option<Self::Item> {
        self.label += 1;
        self.istart = self.iend;
        self.iend = *self.graph.node_index.get(self.label)?;
        Some(&self.graph.neighbours[self.istart..self.iend])
    }
}

impl Graph {
    /// Returns the raw indices of the compressed sparse [`Graph`] struct
    ///
    /// See [`from_raw_parts`](Self::from_raw_parts) for information.
    pub fn get_raw_node_index(&self) -> &[usize] {
        &self.node_index
    }

    /// Returns the raw neighbour list of the compressed sparse [`Graph`] struct
    ///
    /// See [`from_raw_parts`](Self::from_raw_parts) for information.
    pub fn get_raw_neighbours(&self) -> &[usize] {
        &self.neighbours
    }
}

impl Graph {
    /// Checks if all neighbours are also in the graph
    pub fn check_valid_neighbours(&self) -> bool {
        self.neighbours
            .iter()
            .all(|&label| label < self.num_nodes())
    }

    /// Checks if each link in the graph goes in both directions
    pub fn check_bijectivity(&self) -> bool {
        for (label, nbrs) in self.iter().enumerate() {
            for &nbr in nbrs {
                let Some(back_nbrs) = self.get_neighbours_checked(nbr) else {
                    debug!("Graph contains neighbour {nbr} that are not in the graph.");
                    return false;
                };
                if !back_nbrs.contains(&label) {
                    debug!("Link {label}, {nbr} does not go in both directions.");
                    return false;
                }
            }
        }

        true
    }
}

impl Graph {
    /// Generate a two-dimensional cubic or square graph with edge length `size`,
    /// such that there are `size * size` nodes and each node has 4 neighbours.
    pub fn cubic2d(size: usize) -> Self {
        let n = size;
        let nodes = iproduct!(0..n, 0..n).map(|(i, j)| {
            let left = i * n + (j + n - 1) % n;
            let right = i * n + (j + 1) % n;
            let up = ((i + 1) % n) * n + j;
            let down = ((i + n - 1) % n) * n + j;
            [left, up, right, down].into_iter()
        });
        Self::from_neighbour_iter(nodes)
    }

    pub fn cubic3d(size: usize) -> Self {
        let n = size;
        let nodes = iproduct!(0..n, 0..n, 0..n).map(|(i, j, k)| {
            // The node itself has label n(n*i + j) + k
            [
                n * (n * i + j) + ((k + 1) % n),
                n * (n * i + j) + ((k + n - 1) % n),
                n * (n * i + ((j + 1) % n)) + k,
                n * (n * i + ((j + n - 1) % n)) + k,
                n * (n * ((i + 1) % n) + j) + k,
                n * (n * ((i + n - 1) % n) + j) + k,
            ]
            .into_iter()
        });
        Self::from_neighbour_iter(nodes)
    }

    pub fn cubic4d(size: usize) -> Self {
        let n = size;
        let nodes = iproduct!(0..n, 0..n, 0..n, 0..n).map(|(i, j, k, l)| {
            // The node itself has label n(n(n*i + j) + k) + l
            [
                n * (n * (n * i + j) + k) + ((l + 1) % n),
                n * (n * (n * i + j) + k) + ((l + n - 1) % n),
                n * (n * (n * i + j) + ((k + 1) % n)) + l,
                n * (n * (n * i + j) + ((k + n - 1) % n)) + l,
                n * (n * (n * i + ((j + 1) % n)) + k) + l,
                n * (n * (n * i + ((j + n - 1) % n)) + k) + l,
                n * (n * (n * ((i + 1) % n) + j) + k) + l,
                n * (n * (n * ((i + n - 1) % n) + j) + k) + l,
            ]
            .into_iter()
        });
        Self::from_neighbour_iter(nodes)
    }

    /// Create a [`Graph`] of a hexagonal/triangular lattice with side lengths `size`
    ///
    /// This means the total number of nodes in this graph is `size*size`.
    pub fn hexagonal(size: usize) -> Self {
        let n = size;
        let nodes = iproduct!(0..n, 0..n).map(|(i, j)| {
            // The node itself has label: n*i + j
            [
                i * n + (j + n - 1) % n,
                i * n + (j + 1) % n,
                ((i + 1) % n) * n + j,
                ((i + n - 1) % n) * n + j,
                ((i + 1) % n) * n + (j + n - 1) % n,
                ((i + n - 1) % n) * n + (j + 1) % n,
            ]
            .into_iter()
        });
        Self::from_neighbour_iter(nodes)
    }
}

/// Graph with a field on nodes
#[derive(Debug, Clone, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
pub struct FieldGraph<T> {
    #[serde(flatten)]
    graph: Graph,
    field: Vec<T>,
}

impl<T> FieldGraph<T> {
    /// Create a graph with a field on the nodes
    ///
    /// # Panic
    /// This requires the field to define a value for exactly every node, if
    /// the number of field values and nodes do not match this panics.
    pub fn new(graph: Graph, field: Vec<T>) -> Self {
        if graph.num_nodes() != field.len() {
            panic!("Field in not given for exactly all nodes")
        }
        Self { graph, field }
    }
}

impl<T> FieldGraph<T> {
    pub fn get_field(&self) -> &[T] {
        &self.field
    }

    pub fn get_mut_field(&mut self) -> &mut [T] {
        &mut self.field
    }

    pub fn get_field_value_checked(&self, label: usize) -> Option<&T> {
        self.field.get(label)
    }

    pub fn get_mut_field_value_checked(&mut self, label: usize) -> Option<&mut T> {
        self.field.get_mut(label)
    }

    pub fn graph(&self) -> &Graph {
        &self.graph
    }

    pub fn take_graph(self) -> Graph {
        self.graph
    }
}

impl std::fmt::Display for Graph {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Graph({})",
            self.iter()
                .enumerate()
                .map(|(i, nbrs)| {
                    format!(
                        "{}: {}",
                        i,
                        nbrs.iter().map(|label| label.to_string()).join(" ")
                    )
                })
                .join(", ")
        )
    }
}

#[cfg(test)]
mod tests {
    use std::vec;

    use super::*;

    /// Example graph
    /// 0--1--2  3
    ///    | /
    ///    4
    fn example_graph() -> Graph {
        let node_index = vec![0, 1, 4, 6, 6, 8];
        let neighbours = vec![1, 0, 4, 2, 1, 4, 2, 1];
        Graph::from_raw_parts(node_index, neighbours).unwrap()
    }

    #[test]
    fn test_empty_graph() {
        let graph = Graph::default();
        assert!(graph.is_empty())
    }

    #[test]
    fn test_example_graph() {
        let graph = example_graph();

        assert_eq!(graph.num_nodes(), 5);
        assert_eq!(graph.num_links(), 4);

        assert_eq!(graph.get_degree(0), 1);
        assert_eq!(graph.get_degree(1), 3);
        assert_eq!(graph.get_degree(2), 2);
        assert_eq!(graph.get_degree(3), 0);
        assert_eq!(graph.get_degree(4), 2);

        assert_eq!(graph.get_neighbours(0), [1]);
        assert_eq!(graph.get_neighbours(1), [0, 4, 2]);
        assert_eq!(graph.get_neighbours(2), [1, 4]);
        assert_eq!(graph.get_neighbours(3), Vec::<usize>::new().as_slice());
        assert_eq!(graph.get_neighbours(4), [2, 1]);
    }

    #[test]
    fn test_graph_creation() {
        let graph = example_graph();

        let adjacency = [vec![1], vec![0, 4, 2], vec![1, 4], vec![], vec![2, 1]];
        let graph_from_slices = Graph::from_neighbour_slice(adjacency.iter());
        eprintln!("{:?}", graph_from_slices);
        assert_eq!(graph, graph_from_slices);
        let graph_from_iter =
            Graph::from_neighbour_iter(adjacency.into_iter().map(|nbrs| nbrs.into_iter()));
        assert_eq!(graph, graph_from_iter);
    }

    #[test]
    fn test_graph_iter() {
        let graph = example_graph();
        let mut graph_iter = graph.iter();
        assert_eq!(graph_iter.next(), Some(&[1][..]));
        assert_eq!(graph_iter.next(), Some(&[0, 4, 2][..]));
        assert_eq!(graph_iter.next(), Some(&[1, 4][..]));
        assert_eq!(graph_iter.next(), Some(&[][..]));
        assert_eq!(graph_iter.next(), Some(&[2, 1][..]));
        assert_eq!(graph_iter.next(), None);
    }

    #[test]
    fn test_graph_with_field() {
        let graph = example_graph();
        let field = vec![3, 5, 1, 8, 9];
        let graph = FieldGraph::new(graph, field);

        let graph_str = serde_json::to_string(&graph).unwrap();
        let json_graph = serde_json::from_str(&graph_str).unwrap();

        assert_eq!(graph, json_graph);
    }
}