hypergraphx 0.0.5

use crate::{
    HypergraphErrors,
    core::{
        Graph,
        multiplex::{
            MultiplexBase, MultiplexNode, WildcardEdgeType,
            directive::{MultiplexEdgeType, UndirectedMultiplexEdge},
            weak_layers::WeakULayer,
        },
    },
    traits::*,
};
use hashbrown::{HashMap, HashSet};
use std::{
    any::{TypeId, type_name_of_val},
    hash::Hash,
    rc::Rc,
};

use itertools::Itertools;

#[derive(Clone, PartialEq, Eq)]
pub struct UndirectedLayer<N, L> {
    pub(crate) edges: Vec<UndirectedMultiplexEdge>,
    pub(crate) weight: L,
    pub(crate) edge_type: TypeId,
    pub(crate) nodes: Rc<MultiplexBase<N>>,

    pub(crate) incidence: HashMap<usize, Vec<usize>>,
}

impl<N, L> UndirectedLayer<N, L>
where
    L: Clone,
    N: Clone,
{
    pub fn new<E: 'static>(weight: L, nodes: Rc<MultiplexBase<N>>) -> Self {
        Self {
            edges: Vec::new(),
            weight,
            edge_type: TypeId::of::<E>(),
            nodes,
            incidence: HashMap::new(),
        }
    }

    pub fn get_neighbours(&self, node_index: usize) -> Option<HashSet<&usize>> {
        if node_index >= self.nodes.nodes.len() {
            return None; // Node index out of bounds
        }
        let mut neighbours = HashSet::new();
        for edge_index in self.incidence[&node_index].iter() {
            if let Some(edge) = self.edges.get(*edge_index) {
                neighbours.extend(edge.nodes.iter());
            }
        }
        neighbours.remove(&node_index); // Remove self from neighbours
        Some(neighbours)
    }

    pub fn add_edge<E: MultiplexEdgeType>(
        &mut self,
        weight: E,
        nodes: Vec<usize>,
    ) -> Result<(), HypergraphErrors> {
        let wtyp = TypeId::of::<E>();
        if self.edge_type == TypeId::of::<WildcardEdgeType>() {
            self.edge_type = wtyp; // Set the edge type if it's a wildcard
        } else if wtyp != self.edge_type {
            return Err(HypergraphErrors::EdgeTypeMismatch {
                expected: self.edges[0].weight.type_name(),
                found: type_name_of_val(&weight),
            });
        }
        if nodes.len() < 2 {
            return Err(HypergraphErrors::EdgeTooSmall);
        }

        for node in nodes.iter() {
            self.incidence
                .entry(*node)
                .or_insert_with(Vec::new)
                .push(self.edges.len());
        }

        let wbox = Box::new(weight);
        let edge = UndirectedMultiplexEdge {
            weight: wbox,
            weight_type: wtyp,
            nodes,
        };

        self.edges.push(edge);

        Ok(())
    }

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

    pub fn remove_edge(&mut self, edge_index: usize) -> Option<UndirectedMultiplexEdge> {
        if edge_index >= self.edges.len() {
            return None;
        }

        let removed_edge = self.edges.swap_remove(edge_index);

        // Remove this edge from the nodes it connects
        for &node_index in &removed_edge.nodes {
            self.incidence
                .entry(node_index)
                .and_modify(|edges| edges.retain(|&e| e != 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() {
                self.incidence.entry(node_index).and_modify(|edges| {
                    for e in edges.iter_mut() {
                        if *e == l {
                            *e = edge_index; // Update moved edge's index
                        }
                    }
                });
            }
        }

        Some(removed_edge)
    }

    pub(crate) fn restrict_to(&self, node_map: &Vec<Option<usize>>) -> Self {
        let mut edge_map = vec![None; self.edges.len()];
        let new_edges = self
            .edges
            .iter()
            .enumerate()
            .map(|(i, x)| (i, x.remap(node_map)))
            .enumerate()
            .filter_map(|(j, (i, x))| {
                if let Some(ref _u) = x {
                    edge_map[i] = Some(j);
                    x
                } else {
                    None
                }
            })
            .collect();
        let mut new_incidence = HashMap::new();
        for (n, es) in self.incidence.iter() {
            if let Some(new_n) = node_map[*n] {
                let new_edges: Vec<usize> = es.iter().filter_map(|&e| edge_map[e]).collect();
                new_incidence.insert(new_n, new_edges);
            }
        }
        Self {
            edges: new_edges,
            weight: self.weight.clone(),
            edge_type: self.edge_type,
            nodes: self.nodes.clone(),
            incidence: new_incidence,
        }
    }

    pub(crate) fn weaken(self) -> WeakULayer<L> {
        WeakULayer {
            weight: self.weight,
            edge_type: self.edge_type,
            edges: self.edges,
            incidence: self.incidence,
        }
    }
}

impl<'a, N: 'a, L> GraphBasics<'a> for UndirectedLayer<N, L>
where
    N: Clone + Eq + Hash,
{
    type NodeRef = &'a MultiplexNode<N>;

    type EdgeRef = &'a UndirectedMultiplexEdge;
    type EdgeIndex = usize;

    type NodeIndex = usize;

    fn nodes(&'a self) -> impl Iterator<Item = Self::NodeRef> {
        self.nodes.iter()
    }

    fn edges(&'a self) -> impl Iterator<Item = Self::EdgeRef> {
        self.edges.iter()
    }

    fn node_count(&'a self) -> usize {
        Rc::as_ref(&self.nodes).nodes.len()
    }

    fn edge_count(&'a self) -> usize {
        self.edges.len()
    }

    fn is_directed(&self) -> bool {
        false
    }

    fn node(&'a self, node_index: Self::NodeIndex) -> Option<Self::NodeRef> {
        self.nodes.get(node_index).map(|node| &*node)
    }

    fn edge(&'a self, edge_index: Self::EdgeIndex) -> Option<Self::EdgeRef> {
        self.edges.get(edge_index).map(|edge| &*edge)
    }

    fn node_iter(
        &'a self,
        node_index: impl Iterator<Item = Self::NodeIndex>,
    ) -> impl Iterator<Item = Option<Self::NodeRef>> {
        node_index.map(move |i| self.node(i))
    }

    fn edge_iter(
        &'a self,
        edge_index: impl Iterator<Item = Self::EdgeIndex>,
    ) -> impl Iterator<Item = Option<Self::EdgeRef>> {
        edge_index.map(move |i| self.edge(i))
    }
}

impl<'a, N: 'a, L> GraphProperties<'a> for UndirectedLayer<N, L>
where
    N: Eq + Hash + Clone,
    L: Clone,
{
    fn incident_edges(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.incidence
                .get(&node_index)?
                .iter()
                // .filter_map(|&edge_index| self.edges.get(edge_index).map(|edge| (edge_index, edge)))
                .map(|&edge_index| edge_index)
                .collect(),
        )
    }

    fn neighbours(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        let mut out: HashSet<usize> = self
            .incidence
            .get(&node_index)?
            .iter()
            .filter_map(|&edge_index| self.edges.get(edge_index))
            .flat_map(|edge| edge.nodes.iter().map(|n| *n))
            .collect();

        out.remove(&node_index);
        Some(out)
    }

    fn degree(&self, node_index: usize) -> Option<usize> {
        Some(self.incidence.get(&node_index)?.len())
    }

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

        for i in 0..self.node_count() {
            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 = Self;

    fn component(&self, node_index: usize) -> Option<Vec<usize>> {
        let mut visited = vec![false; self.node_count()];
        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)
    }

    /// THIS FUNCTION DOES NOT WORK.
    ///
    /// I MIGHT REMOVE IT FROM THE TRAIT.
    fn extract_component(&self, node_index: usize) -> Option<Self::Subgraph> {
        unreachable!()
    }

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

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

    fn degree_by_order(
        &'a self,
        node_index: <Self as GraphBasics<'a>>::NodeIndex,
        order: usize,
    ) -> Option<usize> {
        // self.nodes.get(node_index).map(|node| self.incidence[&node_index].iter().filter(|e| self.edges[**e].nodes.len() == order).count())
        Some(
            self.incidence
                .get(&node_index)?
                .iter()
                .filter(|e| self.edges[**e].nodes.len() == order)
                .count(),
        )
    }
}

impl<'a, N, L> HypergraphBasics<'a> for UndirectedLayer<N, L>
where
    N: Clone + Eq + Hash + 'a,
    L: Clone + Eq + Hash + 'a,
{
    fn uniform(&self) -> bool {
        if self.edges.is_empty() {
            return true;
        }
        let first_len = self.edges[0].nodes.len();
        self.edges.iter().all(|e| e.nodes.len() == first_len)
    }

    type DualType = ();

    fn dual(&self) -> Self::DualType {
        ()
        // let mut out = UndirectedHyperGraph::new();
        // assert!(self.edge_type == TypeId::of::<E>());

        // out.nodes = self
        //     .edges
        //     .iter()
        //     .cloned()
        //     .map(|e| {
        //         let weight = (e.weight as Box<dyn Any + 'static>).downcast::<E>();
        //         let weight = (*weight.unwrap()).clone();

        //         UNode {
        //             weight,
        //             edges: e.nodes.clone(), // Will be filled later
        //         }
        //     })
        //     .collect();

        // out.edges = self
        //     .nodes()
        //     .enumerate()
        //     .map(|(i, n)| UEdge {
        //         weight: n.weight.clone(),
        //         nodes: self.incidence[&i].clone(),
        //     })
        //     .collect();
        // (self.weight.clone(), out)
    }
}

impl<'a, N: 'a, L: 'a> HypergraphProperties<'a, N, ()> for UndirectedLayer<N, L>
where
    N: Clone + Eq + Hash,
    L: Clone + Eq + Hash,
{
    fn order(&self, edge_index: usize) -> Option<usize> {
        self.edges.get(edge_index).map(|edge| edge.nodes.len())
    }

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

        out
    }
}