use std::fmt::Display;

use itertools::{Itertools, iproduct};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Node {
    neighbours: Vec<usize>,
}

impl Node {
    #[inline]
    pub fn new(neighbours: Vec<usize>) -> Self {
        Self { neighbours }
    }

    #[inline]
    pub fn get_neighbours(&self) -> &[usize] {
        &self.neighbours
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.neighbours.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.neighbours.is_empty()
    }

    #[inline]
    pub fn degree(&self) -> usize {
        self.len()
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Graph {
    nodes: Vec<Node>,
}

impl Graph {
    #[inline]
    pub fn len(&self) -> usize {
        self.nodes.len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.nodes.is_empty()
    }

    #[inline]
    pub fn get_node(&self, label: usize) -> &Node {
        &self.nodes[label]
    }

    #[inline]
    pub fn nodes(&self) -> &[Node] {
        &self.nodes
    }

    #[inline]
    pub fn get_neighbours(&self, label: usize) -> &[usize] {
        &self.get_node(label).neighbours
    }
}

impl Graph {
    /// Create a new [`Graph`] for given `nodes`
    ///
    /// Requires `nodes` to only point to nodes that are part of `nodes`, i.e.
    /// all nodes in the graph must be part of the list.
    #[inline]
    pub fn new(nodes: Vec<Node>) -> Self {
        Self { nodes }
    }

    /// Create a new [`Graph`] for given adjacency 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`.
    pub fn new_from_iter<I, J>(adjacency: I) -> Self
    where
        I: Iterator<Item = J>,
        J: Iterator<Item = usize>,
    {
        let nodes = adjacency.map(|nbrs| Node::new(nbrs.collect())).collect();
        Self { nodes }
    }
}

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;
                Node {
                    neighbours: vec![left, up, right, down],
                }
            })
            .collect();

        Self { 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
                let neighbours = vec![
                    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,
                ];
                Node { neighbours }
            })
            .collect();

        Self { 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
                let neighbours = vec![
                    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,
                ];
                Node { neighbours }
            })
            .collect();

        Self { 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
                let neighbours = vec![
                    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,
                ];
                Node { neighbours }
            })
            .collect();

        Self { nodes }
    }
}

impl Graph {
    /// Checks if each edge goes in both directions
    pub fn is_bidirectional(&self) -> bool {
        for (label, node) in self.nodes.iter().enumerate() {
            for &nbr in &node.neighbours {
                if !self.get_neighbours(nbr).contains(&label) {
                    return false;
                }
            }
        }
        true
    }
}

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

impl From<Graph> for super::compressed_graph::Graph {
    fn from(value: Graph) -> Self {
        (&value).into()
    }
}

impl From<&Graph> for super::compressed_graph::Graph {
    fn from(value: &Graph) -> Self {
        Self::from_neighbour_slice(value.nodes.iter().map(|node| node.get_neighbours()))
    }
}
