hypergraphx 0.0.5

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

use hashbrown::HashSet;

use super::TemporalBehaviour;
use crate::prelude::*;

/// A directed temporal hypergraph. Some edges be, some edges don't be.
#[derive(Debug, Clone)]
pub struct TemporalHypergraph<N, E> {
    inner: DirectedHypergraph<N, E>,
    temporal_map: Vec<TemporalBehaviour>,
    pub ts: usize,
}

impl_graph_basics_wrapper!(
    TemporalHypergraph<N, E>,
    DirectedHypergraph<N, E>,
    true
);

impl<N, E> TemporalHypergraph<N, E>
where
    N: Clone + Eq + Hash,
    E: Clone + Eq + Hash,
{
    pub fn new() -> Self {
        Self {
            inner: DirectedHypergraph::new(),
            temporal_map: Vec::new(),
            ts: 0,
        }
    }

    pub fn from_inner(inner: DirectedHypergraph<N, E>) -> Self {
        Self {
            temporal_map: vec![TemporalBehaviour::Permanent; inner.edge_count()],
            inner,
            ts: 0,
        }
    }

    pub fn into_inner(self) -> DirectedHypergraph<N, E> {
        self.inner
    }

    pub fn update_edge_behaviour<'a>(
        &'a mut self,
        edge: <Self as GraphBasics<'a>>::EdgeIndex,
        behaviour: TemporalBehaviour,
    ) -> bool {
        if edge < self.temporal_map.len() {
            self.temporal_map[edge] = behaviour;
            true
        } else {
            false
        }
    }

    pub fn add_node(&mut self, node: N) -> usize {
        self.inner.add_node(node)
    }
    pub fn add_edge(
        &mut self,
        edge: E,
        behaviour: TemporalBehaviour,
        source: Vec<usize>,
        target: Vec<usize>,
    ) -> Result<(), HypergraphErrors> {
        let edge_index = self.inner.add_edge(edge, source, target)?;
        if edge_index >= self.temporal_map.len() {
            self.temporal_map.push(behaviour);
        }
        Ok(())
    }
    pub fn add_nodes(&mut self, nodes: impl Iterator<Item = N>) {
        self.inner.add_nodes(nodes)
    }
    pub fn add_edges(
        &mut self,
        edges: impl Iterator<Item = (E, Vec<usize>, Vec<usize>)>,
        behs: impl Iterator<Item = TemporalBehaviour>,
    ) -> Result<(), HypergraphErrors> {
        self.inner.add_edges(edges)?;
        self.temporal_map.extend(behs);
        Ok(())
    }

    pub fn remove_node(&mut self, node_index: usize) -> Option<DirectedNode<N>> {
        if node_index < self.node_count() {
            let removed_node = &self.inner.nodes[node_index];
            let mut out_edges = vec![];
            // Remove edges connected to this node
            for &edge_index in &removed_node.in_edges {
                if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                    edge.target.retain(|&n| n != node_index);
                    if edge.target.len() < 1 {
                        // If the edge has less than 2 nodes, remove it
                        out_edges.push(edge_index);
                    }
                }
            }

            for &edge_index in &removed_node.out_edges {
                if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                    edge.source.retain(|&n| n != node_index);
                    if edge.source.len() < 1 {
                        // If the edge has less than 2 nodes, remove it
                        out_edges.push(edge_index);
                    }
                }
            }

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

            let l = self.node_count();
            if l > node_index {
                let moved_node = &mut self.inner.nodes[node_index];
                for &edge_index in moved_node.out_edges.iter() {
                    if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                        edge.source.iter_mut().for_each(|n| {
                            if *n == l {
                                *n = node_index; // Update moved node's index
                            }
                        });
                    }
                }
                for &edge_index in moved_node.in_edges.iter() {
                    if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                        edge.target.iter_mut().for_each(|n| {
                            if *n == l {
                                *n = node_index; // Update moved node's index
                            }
                        });
                    }
                }
            }
            Some(removed_node)
        } else {
            None
        }
    }

    pub fn remove_nodes(&mut self, node_indices: Vec<usize>) -> Vec<DirectedNode<N>> {
        for &node_index in node_indices.iter() {
            let removed_node = &self.inner.nodes[node_index];
            let mut out_edges = vec![];
            // Remove edges connected to this node
            for &edge_index in &removed_node.in_edges {
                if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                    edge.target.retain(|&n| n != node_index);
                    if edge.target.len() < 1 {
                        // If the edge has less than 2 nodes, remove it
                        out_edges.push(edge_index);
                    }
                }
            }

            for &edge_index in &removed_node.out_edges {
                if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                    edge.source.retain(|&n| n != node_index);
                    if edge.source.len() < 1 {
                        // If the edge has less than 2 nodes, remove it
                        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.node_count() - count - 1;
            if l > node_index {
                let moved_node = &mut self.inner.nodes[l];
                for &edge_index in moved_node.in_edges.iter() {
                    if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                        edge.target.iter_mut().for_each(|n| {
                            if *n == l {
                                *n = node_index; // Update moved node's index
                            }
                        });
                    }
                }

                for &edge_index in moved_node.out_edges.iter() {
                    if let Some(edge) = self.inner.edges.get_mut(edge_index) {
                        edge.source.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.node_count() {
                    Some(self.inner.nodes.swap_remove(node_index))
                } else {
                    None
                }
            })
            .collect()
    }

    pub fn remove_edge(&mut self, edge_index: usize) -> Option<DirectedEdge<E>> {
        if edge_index < self.edge_count() {
            let removed_edge = &self.inner.edges[edge_index];
            // Remove this edge from the nodes it connects
            for &node_index in &removed_edge.source {
                if let Some(node) = self.inner.nodes.get_mut(node_index) {
                    node.out_edges.retain(|&e| e != edge_index);
                }
            }

            for &node_index in &removed_edge.target {
                if let Some(node) = self.inner.nodes.get_mut(node_index) {
                    node.in_edges.retain(|&e| e != edge_index);
                }
            }

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

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

                for &node_index in moved_edge.target.iter() {
                    if let Some(node) = self.inner.nodes.get_mut(node_index) {
                        node.in_edges.iter_mut().for_each(|e| {
                            if *e == l {
                                *e = edge_index; // Update moved edge's index
                            }
                        });
                    }
                }
            }

            Some(removed_edge)
        } else {
            None
        }
    }

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

                for &node_index in &removed_edge.target {
                    if let Some(node) = self.inner.nodes.get_mut(node_index) {
                        node.in_edges.retain(|&e| e != edge_index);
                    }
                }
            }
        }

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

                for &node_index in moved_edge.target.iter() {
                    if let Some(node) = self.inner.nodes.get_mut(node_index) {
                        node.in_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.edge_count() {
                    Some(self.inner.edges.swap_remove(edge_index))
                } else {
                    None
                }
            })
            .collect()
    }

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

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

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

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

impl<'a, N, E> TemporalProperties<'a> for TemporalHypergraph<N, E>
where
    N: Clone + Eq + Hash + 'a,
    E: Clone + Eq + Hash + 'a,
{
    fn get_time(&self) -> usize {
        self.ts
    }

    fn get_time_mut(&mut self) -> &mut usize {
        &mut self.ts
    }

    type Inner = DirectedHypergraph<N, E>;

    fn snapshot(&'a self) -> &'a Self::Inner {
        &self.inner
    }

    fn temporal_behaviour(
        &self,
        edge_index: <Self as GraphBasics<'a>>::EdgeIndex,
    ) -> Option<TemporalBehaviour> {
        self.temporal_map.get(edge_index).map(|x| *x)
    }
}