hypergraphx 0.0.5

use hashbrown::{HashMap, HashSet};
use nalgebra::{Const, CsMatrix, DMatrix, Dyn, Matrix};
use std::{
    any::{TypeId, type_name_of_val},
    hash::Hash,
    rc::Rc,
};

use crate::{
    HypergraphErrors,
    core::{
        DiGraph,
        directed::{Hypergraph, Node, TarjanState},
        multiplex::{
            MultiplexBase, MultiplexNode, WildcardEdgeType,
            directive::{DirectedMultiplexEdge, MultiplexEdgeType},
            weak_layers::WeakDLayer,
        },
    },
    traits::*,
};

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

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

impl<N, L> DirectedLayer<N, L> {
    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_out_neighbours(&self, node_index: usize) -> Option<HashSet<&usize>> {
        let mut neighbours = HashSet::new();

        for edge_index in self.incidence.get(&node_index)?.1.iter() {
            if let Some(edge) = self.edges.get(*edge_index) {
                neighbours.extend(edge.target.iter());
            }
        }
        Some(neighbours)
    }

    pub fn get_in_neighbours(&self, node_index: usize) -> Option<HashSet<&usize>> {
        let mut neighbours = HashSet::new();

        for edge_index in self.incidence.get(&node_index)?.0.iter() {
            if let Some(edge) = self.edges.get(*edge_index) {
                neighbours.extend(edge.source.iter());
            }
        }
        Some(neighbours)
    }

    pub fn add_edge<E: MultiplexEdgeType>(
        &mut self,
        weight: E,
        source: Vec<usize>,
        target: 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 source.is_empty() || target.is_empty() {
            return Err(HypergraphErrors::EdgeTooSmall);
        }

        for node in source.iter() {
            self.incidence
                .entry(*node)
                .or_insert_with(|| (Vec::new(), Vec::new()))
                .1
                .push(self.edges.len());
        }
        for node in target.iter() {
            self.incidence
                .entry(*node)
                .or_insert_with(|| (Vec::new(), Vec::new()))
                .0
                .push(self.edges.len());
        }

        let wbox = Box::new(weight);
        let edge = DirectedMultiplexEdge {
            weight: wbox,
            weight_type: wtyp,
            source,
            target,
        };

        self.edges.push(edge);

        Ok(())
    }

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

    pub fn remove_edge(&mut self, edge_index: usize) -> Option<DirectedMultiplexEdge> {
        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.source {
            self.incidence
                .entry(node_index)
                .and_modify(|edges| edges.0.retain(|&e| e != edge_index));
        }

        for &node_index in &removed_edge.target {
            self.incidence
                .entry(node_index)
                .and_modify(|edges| edges.1.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.source.iter() {
                self.incidence.entry(node_index).and_modify(|edges| {
                    for e in edges.1.iter_mut() {
                        if *e == l {
                            *e = edge_index; // Update moved edge's index
                        }
                    }
                });
            }
            for &node_index in moved_edge.target.iter() {
                self.incidence.entry(node_index).and_modify(|edges| {
                    for e in edges.0.iter_mut() {
                        if *e == l {
                            *e = edge_index; // Update moved edge's index
                        }
                    }
                });
            }
        }

        Some(removed_edge)
    }

    fn tarjan_inner(
        &self,
        curr_idx: &mut usize,
        curr: usize,
        aux: &mut Vec<TarjanState>,
        stack: &mut Vec<usize>,
    ) -> Option<Vec<usize>>
    where
        // E: Clone + Eq + Hash + 'static,
        N: Clone + Eq + Hash,
    {
        let neighbours = {
            let state = &mut aux[curr];

            state.index = Some(*curr_idx);
            state.lowlink = Some(*curr_idx);
            *curr_idx += 1;
            state.on_stack = true;

            let neighbours = self.get_out_neighbours(curr);

            stack.push(curr);
            neighbours
        }?;

        for n in neighbours {
            let nb = aux[*n];
            if nb.index.is_none() {
                // Successor has not been visited yet
                self.tarjan_inner(curr_idx, *n, aux, stack);
                let state = &mut aux[curr];
                state.lowlink = Some(state.lowlink.unwrap().min(nb.lowlink.unwrap()));
            } else if nb.on_stack {
                // Successor is on the stack
                let state = &mut aux[curr];
                state.lowlink = Some(state.lowlink.unwrap().min(nb.index.unwrap()));
            }
        }

        let state = aux[curr];
        if state.lowlink == state.index {
            // Found a strongly connected component
            let mut component = Vec::new();
            while let Some(node_index) = stack.pop() {
                let node_state = &mut aux[node_index];
                node_state.on_stack = false;
                component.push(node_index);
                if node_index == curr {
                    break;
                }
            }

            return Some(component);
        }

        None
    }

    pub(crate) fn restrict_to(&self, node_map: &Vec<Option<usize>>) -> Self
    where
        L: Clone,
    {
        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_in: Vec<usize> = es.0.iter().filter_map(|&e| edge_map[e]).collect();
                let new_out: Vec<usize> = es.1.iter().filter_map(|&e| edge_map[e]).collect();
                new_incidence.insert(new_n, (new_in, new_out));
            }
        }
        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) -> WeakDLayer<L> {
        WeakDLayer {
            weight: self.weight,
            edge_type: self.edge_type,
            edges: self.edges,
            incidence: self.incidence,
        }
    }

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

        let mut new_nodes = vec![];

        for &node_index in node_indices.iter() {
            if let Some(node) = self.nodes.get(node_index) {
                let new_index = new_nodes.len();
                // TODO: This clone is probably preventable.
                new_nodes.push(Node {
                    weight: node.weight.clone(),
                    in_edges: self.incidence[&node_index].0.clone(),
                    out_edges: self.incidence[&node_index].1.clone(),
                });
                node_map[node_index] = Some(new_index);
            }
        }

        let mut u = node_indices
            .iter()
            .flat_map(|i| {
                self.incidence[i]
                    .0
                    .iter()
                    .chain(self.incidence[i].1.iter())
                    .map(|x| (*x, &self.edges[*x]))
            })
            .collect::<Vec<_>>();
        u.retain(|(_i, e)| {
            !e.source.is_empty()
                && e.source.iter().all(|&n| node_map[n].is_some())
                && !e.target.is_empty()
                && e.target.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.in_edges = node
                .in_edges
                .iter()
                .filter_map(|&e| edge_map.get(&e))
                .cloned()
                .collect();

            node.out_edges = node
                .out_edges
                .iter()
                .filter_map(|&e| edge_map.get(&e))
                .cloned()
                .collect();
        }

        let mut out = Hypergraph::new();
        out.nodes = new_nodes;
        for (_, edge) in &u {
            let new_source: Vec<usize> = edge.source.iter().filter_map(|&n| node_map[n]).collect();
            let new_target: Vec<usize> = edge.target.iter().filter_map(|&n| node_map[n]).collect();
            if !new_source.is_empty() && !new_target.is_empty() {
                let _ = out.add_edge((), new_source, new_target);
            }
        }

        out
    }
}

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

    type EdgeRef = &'a DirectedMultiplexEdge;
    type NodeIndex = usize;

    type EdgeIndex = 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 {
        true
    }

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

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

    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>> {
        // todo!()
        edge_index.map(move |i| self.edge(i))
    }
}

impl<'a, N: 'a, L> DiGraphProperties<'a> for DirectedLayer<N, L>
where
    N: Eq + Hash + Clone,
{
    fn source(
        &'a self,
        edge_index: <Self as GraphBasics<'a>>::EdgeIndex,
    ) -> Option<hashbrown::HashSet<usize>> {
        let edge = self.edges.get(edge_index)?;
        let sources = edge
            .source
            .iter()
            // .filter_map(|&n| self.nodes.get(n).map(|node| (n, node)))
            .map(|&n| n)
            .collect::<HashSet<_>>();
        Some(sources)
    }

    fn target(
        &'a self,
        edge_index: <Self as GraphBasics<'a>>::EdgeIndex,
    ) -> Option<hashbrown::HashSet<usize>> {
        let edge = self.edges.get(edge_index)?;
        let targets = edge
            .target
            .iter()
            // .filter_map(|&n| self.nodes.get(n).map(|node| (n, node)))
            .map(|&n| n)
            .collect::<HashSet<_>>();
        Some(targets)
    }

    fn in_edges(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.incidence
                .get(&node_index)?
                .0
                .iter()
                // .filter_map(|&edge_index| self.edges.get(edge_index).map(|edge| (edge_index, edge)))
                .map(|&edge_index| edge_index)
                .collect(),
        )
    }
    fn in_edge_count(&self, node_index: usize) -> Option<usize> {
        Some(self.incidence.get(&node_index)?.0.len())
    }

    fn out_edges(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.incidence
                .get(&node_index)?
                .1
                .iter()
                // .filter_map(|&edge_index| self.edges.get(edge_index).map(|edge| (edge_index, edge)))
                .map(|&edge_index| edge_index)
                .collect(),
        )
    }
    fn out_edge_count(&self, node_index: usize) -> Option<usize> {
        Some(self.incidence.get(&node_index)?.1.len())
    }

    fn in_neighbours(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.in_edges(node_index)?
                .iter()
                .filter_map(|&edge| self.source(edge))
                // .flat_map(|edge| {
                //     edge
                //         .iter()
                //         .filter_map(|&src| self.nodes.get(src).map(|node| (src, node)))
                // })
                .flatten()
                .collect(),
        )
    }

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

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

    fn out_neighbours(&self, node_index: usize) -> Option<hashbrown::HashSet<usize>> {
        Some(
            self.out_edges(node_index)?
                .iter()
                .filter_map(|&edge| self.target(edge))
                // .flat_map(|edge| {
                //     edge
                //         .iter()
                //         .filter_map(|&src| self.nodes.get(src).map(|node| (src, node)))
                // })
                .flatten()
                .collect(),
        )
    }

    fn strongly_connected_components(&self) -> Vec<Vec<usize>> {
        let mut idx = 0;

        let mut aux = vec![
            TarjanState {
                lowlink: None,
                index: None,
                on_stack: false,
            };
            self.node_count()
        ];

        let mut stack = vec![];
        let mut out = vec![];
        for x in 0..aux.len() {
            if aux[x].index.is_none() {
                if let Some(v) = self.tarjan_inner(&mut idx, x, &mut aux, &mut stack) {
                    out.push(v);
                }
            }
        }
        out
    }

    fn strong_component(&self, node_index: usize) -> Option<Vec<usize>> {
        let mut idx = 0;
        let mut aux = vec![
            TarjanState {
                lowlink: None,
                index: None,
                on_stack: false,
            };
            self.node_count()
        ];

        let mut stack = vec![];
        self.tarjan_inner(&mut idx, node_index, &mut aux, &mut stack)
    }

    fn extract_strong_component(&self, node_index: usize) -> Option<Self::Subgraph> {
        let node_indices = self.strong_component(node_index)?;
        Some(self.induced_shgraph(&node_indices))
    }

    fn weakly_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_out_neighbours(node_index).unwrap().iter().cloned());
                        stack.extend(self.get_in_neighbours(node_index).unwrap().iter().cloned());
                    }
                }

                components.push(component);
            }
        }

        components
    }

    fn weak_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_in_neighbours(node_index).unwrap().iter().cloned());
                stack.extend(self.get_out_neighbours(node_index).unwrap().iter().cloned());
            }
        }

        Some(component)
    }

    fn extract_weak_component(&self, node_index: usize) -> Option<Self::Subgraph> {
        let node_indices = self.weak_component(node_index)?;
        Some(self.induced_shgraph(&node_indices))
    }

    fn in_degree(&self, node_index: usize) -> Option<usize> {
        self.incidence
            .get(&node_index)
            .map(|(in_edges, _)| in_edges.len())
    }

    fn out_degree(&self, node_index: usize) -> Option<usize> {
        self.incidence
            .get(&node_index)
            .map(|(_, out_edges)| out_edges.len())
    }

    type Subgraph = Hypergraph<N, ()>;

    fn condense(&self) -> DiGraph<(), ()> {
        todo!()
    }
}

/* impl<'a, N: 'a, L> MatrixRepresentation<'a> for DirectedLayer<N, L>
where
    N: Eq + Hash + Clone,
{
    fn binary_incidence_matrix(&self) -> nalgebra::sparse::CsMatrix<i8> {
        let mut irows = Vec::new();
        let mut icols = Vec::new();
        let mut vals = vec![];

        for (edge_index, edge) in self.edges.iter().enumerate() {
            for &node_index in edge.source.iter() {
                irows.push(node_index);
                icols.push(edge_index);
                vals.push(1);
            }
            for &node_index in edge.target.iter() {
                irows.push(node_index);
                icols.push(edge_index);
                vals.push(-11);
            }
        }

        let matrix =
            CsMatrix::from_triplet(self.node_count(), self.edges.len(), &irows, &icols, &vals);
        matrix
    }

    fn adjacency_matrix(&self) -> DMatrix<usize> {
        let l = self.node_count();
        let mut mat = DMatrix::<usize>::zeros(l, l);

        for (node_index, _) in self.nodes().enumerate() {
            let n = self.get_out_neighbours(node_index).unwrap();
            for &neighbour in n {
                mat[(node_index, neighbour)] += 1;
            }
        }

        mat
    }

    fn dual_adjacency_matrix(&self) -> DMatrix<usize> {
        let l = self.node_count();
        let mut mat = DMatrix::<usize>::zeros(l, l);

        for (node_index, _) in self.nodes().enumerate() {
            let n = self.get_out_neighbours(node_index).unwrap();
            for &neighbour in n {
                mat[(neighbour, node_index)] += 1;
            }
        }

        mat
    }

    fn laplacian_matrix(&'a self) -> DMatrix<usize> {
        let dseq = Matrix::<usize, Dyn, Const<1>, _>::from(self.out_degree_sequence());
        let out = DMatrix::<usize>::from_diagonal(&dseq);

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

    type DualType = ();

    fn dual(&self) -> Self::DualType {
        ()
    }
}

impl<'a, N: 'a, L> DirectedHypergraphProperties<'a, N, ()> for DirectedLayer<N, L>
where
    N: Eq + Hash + Clone,
    L: Clone + 'a,
{
    fn in_order(&self, edge_index: usize) -> Option<usize> {
        self.edges.get(edge_index).map(|edge| edge.source.len())
    }
    fn out_order(&self, edge_index: usize) -> Option<usize> {
        self.edges.get(edge_index).map(|edge| edge.target.len())
    }

    fn digraph_view(&self) -> DiGraph<&N, &()> {
        let mut out = DiGraph::new();
        out.add_nodes(self.nodes.iter().map(|n| &n.weight));
        for e in &self.edges {
            for u in e.source.iter() {
                for v in e.target.iter() {
                    let _ = out.add_edge(&(), [*u], [*v]);
                }
            }
        }

        out
    }
}