hypergraphx 0.0.5

use std::{fmt::Debug, hash::Hash};

use super::{Edge, Hypergraph, Node};
use hashbrown::{HashMap, HashSet};
use itertools::Itertools;

use crate::{HypergraphErrors, core::Graph, impl_graph_basics, impl_weights, traits::*};

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct UniformEdge<E, const ORDER: usize> {
    /// Associated edge data.
    pub weight: E,

    /// Head and Tail of hyperedge
    pub(crate) nodes: [usize; ORDER],
}

impl<E, const ORDER: usize> UniformEdge<E, ORDER> {
    /// Return the source node index.
    pub fn source(&self) -> &[usize] {
        self.nodes.as_slice()
    }

    /// Return the target node index.
    pub fn target(&self) -> &[usize] {
        self.nodes.as_slice()
    }

    pub fn connects(&self, a: usize, b: usize) -> bool {
        self.nodes.contains(&a) && self.nodes.contains(&b)
    }
}

#[derive(Debug, Clone)]
pub struct UniformHypergraph<N, E, const ORDER: usize> {
    pub(crate) nodes: Vec<Node<N>>,
    pub(crate) edges: Vec<UniformEdge<E, ORDER>>,
}

impl<N, E, const ORDER: usize> UniformHypergraph<N, E, ORDER> {
    pub fn new() -> Self {
        Self {
            nodes: Vec::new(),
            edges: Vec::new(),
        }
    }

    pub fn add_node(&mut self, weight: N) -> usize {
        self.nodes.push(Node {
            weight,
            edges: Vec::new(),
        });
        self.nodes.len() - 1
    }

    pub fn add_edge(
        &mut self,
        weight: E,
        node_indices: [usize; ORDER],
    ) -> Result<usize, HypergraphErrors> {
        let edge = UniformEdge {
            weight,
            nodes: node_indices,
        };
        self.edges.push(edge);
        let edge_index = self.edges.len() - 1;

        if self.edges[edge_index]
            .nodes
            .iter()
            .any(|x| self.nodes.get(*x).is_none())
        {
            return Err(HypergraphErrors::Nonexistent);
        }

        for &node_index in &self.edges[edge_index].nodes {
            let node = self.nodes.get_mut(node_index).unwrap();
            node.edges.push(edge_index);
        }

        Ok(edge_index)
    }

    pub fn add_nodes(&mut self, weights: impl Iterator<Item = N>) -> Vec<usize> {
        weights.into_iter().map(|w| self.add_node(w)).collect()
    }

    pub fn add_edges(
        &mut self,
        edges: impl Iterator<Item = (E, [usize; ORDER])>,
    ) -> Result<(), HypergraphErrors> {
        for (weight, node_indices) in edges {
            self.add_edge(weight, node_indices)?;
        }
        Ok(())
    }

    pub fn remove_node(&mut self, node_index: usize) -> Result<Node<N>, HypergraphErrors> {
        if node_index < self.nodes.len() {
            let removed_node = &self.nodes[node_index];
            let mut out_edges = removed_node.edges.clone();
            // Remove edges connected to this node
            // for &edge_index in removed_node.edges.iter() {
            //     out_edges.push(edge_index);
            // }

            out_edges.sort_unstable();
            out_edges.dedup();
            self.remove_edges(out_edges);
            let removed_node = self.nodes.swap_remove(node_index);

            let l = self.nodes.len();
            if l > node_index {
                let moved_node = &mut self.nodes[node_index];
                for &edge_index in moved_node.edges.iter() {
                    if let Some(edge) = self.edges.get_mut(edge_index) {
                        edge.nodes.iter_mut().for_each(|n| {
                            if *n == l {
                                *n = node_index; // Update moved node's index
                            }
                        });
                    }
                }
            }
            Ok(removed_node)
        } else {
            Err(HypergraphErrors::Nonexistent)
        }
    }

    pub fn remove_nodes(&mut self, node_indices: Vec<usize>) -> Vec<Node<N>> {
        for &node_index in node_indices.iter() {
            let removed_node = &self.nodes[node_index];
            let mut out_edges = removed_node.edges.clone();
            // Remove edges connected to this node
            // for &edge_index in removed_node.edges.iter() {
            //     out_edges.push(edge_index);
            // }

            out_edges.sort_unstable();
            out_edges.dedup();
            self.remove_edges(out_edges);
        }

        for (count, &node_index) in node_indices.iter().enumerate() {
            let l = self.nodes.len() - count - 1;
            if l > node_index {
                let moved_node = &mut self.nodes[l];
                for &edge_index in moved_node.edges.iter() {
                    if let Some(edge) = self.edges.get_mut(edge_index) {
                        edge.nodes.iter_mut().for_each(|n| {
                            if *n == l {
                                *n = node_index; // Update moved node's index
                            }
                        });
                    }
                }
            }
        }

        node_indices
            .into_iter()
            .filter_map(|node_index| {
                if node_index < self.nodes.len() {
                    Some(self.nodes.swap_remove(node_index))
                } else {
                    None
                }
            })
            .collect()
    }

    pub fn remove_edge(
        &mut self,
        edge_index: usize,
    ) -> Result<UniformEdge<E, ORDER>, HypergraphErrors> {
        if edge_index < self.edges.len() {
            let removed_edge = &self.edges[edge_index];
            // let out_nodes = vec![];
            // Remove this edge from the nodes it connects
            for &node_index in &removed_edge.nodes {
                if let Some(node) = self.nodes.get_mut(node_index) {
                    node.edges.retain(|&e| e != edge_index);
                }
            }
            let removed_edge = self.edges.swap_remove(edge_index);

            let l = self.edges.len();
            if l > edge_index {
                let moved_edge = &mut self.edges[edge_index];
                for &node_index in moved_edge.nodes.iter() {
                    if let Some(node) = self.nodes.get_mut(node_index) {
                        node.edges.iter_mut().for_each(|e| {
                            if *e == l {
                                *e = edge_index; // Update moved edge's index
                            }
                        });
                    }
                }
            }

            Ok(removed_edge)
        } else {
            Err(HypergraphErrors::Nonexistent)
        }
    }

    pub fn remove_edges(&mut self, edge_indices: Vec<usize>) -> Vec<UniformEdge<E, ORDER>> {
        for &edge_index in edge_indices.iter() {
            if edge_index < self.edges.len() {
                let removed_edge = &self.edges[edge_index];
                // let out_nodes = vec![];
                // Remove this edge from the nodes it connects
                for &node_index in &removed_edge.nodes {
                    if let Some(node) = self.nodes.get_mut(node_index) {
                        node.edges.retain(|&e| e != edge_index);
                    }
                }
            }
        }

        for (count, &edge_index) in edge_indices.iter().rev().enumerate() {
            let l = self.edges.len() - count - 1;
            if l > edge_index {
                let moved_edge = &mut self.edges[l];
                for &node_index in moved_edge.nodes.iter() {
                    if let Some(node) = self.nodes.get_mut(node_index) {
                        node.edges.iter_mut().for_each(|e| {
                            if *e == l {
                                *e = edge_index; // Update moved edge's index
                            }
                        });
                    }
                }
            }
        }

        edge_indices
            .into_iter()
            .rev()
            .filter_map(|edge_index| {
                if edge_index < self.edges.len() {
                    Some(self.edges.swap_remove(edge_index))
                } else {
                    None
                }
            })
            .collect()
    }

    pub fn get_neighbours(&self, node_index: usize) -> Option<HashSet<&usize>> {
        if node_index >= self.nodes.len() {
            return None;
        }
        let mut out = self.nodes[node_index]
            .edges
            .iter()
            .flat_map(|e| {
                if let Some(edge) = self.edges.get(*e) {
                    edge.nodes.iter().collect::<Vec<_>>()
                } else {
                    std::iter::empty().collect()
                }
            })
            .collect::<HashSet<_>>();

        out.remove(&node_index);
        Some(out)
    }
    pub fn get_incident_edges(&self, node_index: usize) -> Option<&Vec<usize>> {
        if node_index >= self.nodes.len() {
            return None;
        }
        Some(&self.nodes[node_index].edges)
    }

    pub fn get_order_if_uniform(&self) -> Option<usize> {
        if self.edges.is_empty() {
            return Some(0);
        }
        let first_len = self.edges[0].nodes.len();
        if self.edges.iter().all(|e| e.nodes.len() == first_len) {
            Some(first_len)
        } else {
            None
        }
    }

    pub fn induced_shgraph(&self, node_indices: &[usize]) -> Hypergraph<N, E>
    where
        E: Clone,
        N: Clone,
    {
        let mut subgraph = Hypergraph::new();
        let mut node_map = vec![None; self.nodes.len()];

        let mut new_nodes = vec![];

        for &node_index in node_indices {
            if let Some(node) = self.nodes.get(node_index) {
                let new_index = new_nodes.len();
                new_nodes.push(node.clone());
                node_map[node_index] = Some(new_index);
            }
        }

        let mut u = new_nodes
            .iter()
            .flat_map(|n| n.edges.iter().map(|x| (*x, &self.edges[*x])))
            .collect::<Vec<_>>();
        u.retain(|(_i, e)| !e.nodes.is_empty() && e.nodes.iter().all(|&n| node_map[n].is_some()));

        let edge_map = u
            .iter()
            .enumerate()
            .map(|(i, (j, _))| (*j, i))
            .collect::<HashMap<_, _>>();

        for node in new_nodes.iter_mut() {
            node.edges = node
                .edges
                .iter()
                .filter_map(|&e| edge_map.get(&e))
                .cloned()
                .collect();
        }

        subgraph.nodes = new_nodes;

        for (_, edge) in &u {
            let new_nodes: Vec<usize> = edge.nodes.iter().filter_map(|&n| node_map[n]).collect();
            if !new_nodes.is_empty() {
                subgraph.add_edge(edge.weight.clone(), new_nodes).unwrap();
            }
        }

        subgraph
    }

    pub fn get_adjacency_tensor(&self) -> ndarray::Array<u8, ndarray::IxDyn> {
        let l = self.nodes.len();
        let mut out = ndarray::Array::zeros(ndarray::IxDyn(&vec![l; ORDER]));

        for edge in self.edges.iter() {
            out[edge.nodes] += 1;
        }

        out
    }
}

impl_graph_basics!(
    UniformHypergraph<N, E, ORDER>,
    &'a Node<N>,
    &'a UniformEdge<E, ORDER>,
    false,
    |const ORDER: usize|
);

impl<'a, N: 'a, E: 'a, const ORDER: usize> GraphProperties<'a> for UniformHypergraph<N, E, ORDER>
where
    N: Clone + Eq + Hash,
    E: Clone + Eq + Hash + Debug,
{
    fn neighbours(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.get_neighbours(node_index)?
                .into_iter()
                .map(|&n| n)
                .collect(),
        )
    }
    fn incident_edges(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.get_incident_edges(node_index)?
                .iter()
                // .filter_map(|&e| self.edges.get(e).map(|edge| (e, edge)))
                .map(|&e| e)
                .collect(),
        )
    }

    fn degree(&self, node_index: usize) -> Option<usize> {
        self.nodes.get(node_index).map(|node| node.edges.len())
    }

    fn neighbour_count(&self, node_index: usize) -> Option<usize> {
        self.nodes.get(node_index).map(|node| {
            node.edges
                .iter()
                .map(|e| self.edges[*e].nodes.len() - 1)
                .sum()
        })
    }

    fn connected_components(&self) -> Vec<Vec<usize>> {
        let mut visited = vec![false; self.nodes.len()];
        let mut components = Vec::new();

        for i in 0..self.nodes.len() {
            if !visited[i] {
                let mut component = Vec::new();
                let mut stack = vec![i];

                while let Some(node_index) = stack.pop() {
                    if !visited[node_index] {
                        visited[node_index] = true;
                        component.push(node_index);
                        stack.extend(self.get_neighbours(node_index).unwrap().iter().cloned());
                    }
                }

                components.push(component);
            }
        }

        components
    }

    type Subgraph = Hypergraph<N, E>;

    fn component(&self, node_index: usize) -> Option<Vec<usize>> {
        if node_index >= self.nodes.len() {
            return None;
        }
        let mut visited = vec![false; self.nodes.len()];
        let mut component = Vec::new();
        let mut stack = vec![node_index];

        while let Some(node_index) = stack.pop() {
            if !visited[node_index] {
                visited[node_index] = true;
                component.push(node_index);
                stack.extend(self.get_neighbours(node_index).unwrap().iter().cloned());
            }
        }

        Some(component)
    }

    fn extract_component(&self, node_index: usize) -> Option<Hypergraph<N, E>> {
        let component_nodes = self.component(node_index)?;
        Some(self.induced_shgraph(&component_nodes))
    }

    fn is_connected(&self) -> bool {
        self.connected_components().len() == 1
    }

    fn contained_nodes(&self, edge_index: usize) -> Option<hashbrown::HashSet<usize>> {
        self.edges
            .get(edge_index)
            .map(|edge| edge.nodes.iter().map(|&n| n).collect())
    }

    fn degree_by_order(
        &'a self,
        node_index: <Self as GraphBasics<'a>>::NodeIndex,
        order: usize,
    ) -> Option<usize> {
        if order != ORDER {
            return None; // Invalid order for uniform hypergraph
        }
        self.degree(node_index)
    }
}

impl<'a, N: 'a, E: 'a, const ORDER: usize> HypergraphBasics<'a> for UniformHypergraph<N, E, ORDER>
where
    E: Clone + Eq + Hash,
    N: Clone + Eq + Hash,
{
    fn uniform(&self) -> bool {
        true
    }

    type DualType = Hypergraph<E, N>;

    fn dual(&self) -> Self::DualType {
        let mut out = Hypergraph::new();

        out.nodes = self
            .edges
            .iter()
            .map(|e| Node {
                weight: e.weight.clone(),
                edges: e.nodes.to_vec(), // Will be filled later
            })
            .collect();

        out.edges = self
            .nodes
            .iter()
            .map(|n| Edge {
                weight: n.weight.clone(),
                nodes: n.edges.clone(),
            })
            .collect();
        out
    }
}

impl<'a, N: 'a, E: 'a, const ORDER: usize> HypergraphProperties<'a, N, E>
    for UniformHypergraph<N, E, ORDER>
where
    N: Clone + Eq + Hash,
    E: Clone + Eq + Hash + Debug,
{
    fn order(&self, edge_index: usize) -> Option<usize> {
        self.edges.get(edge_index).map(|e| e.nodes.len())
    }

    fn graph_view(&self) -> Graph<&N, &E> {
        let mut out = Graph::new();
        out.add_nodes(self.nodes.iter().map(|n| &n.weight));
        for e in &self.edges {
            for (u, v) in e.nodes.iter().tuple_combinations() {
                out.add_edge(&e.weight, [*u, *v]).unwrap();
            }
        }

        out
    }
}

impl_weights!(UniformHypergraph<N, E, ORDER>, |const ORDER: usize|);