use anyhow::{anyhow, Result};
use nalgebra::DMatrix;
use std::collections::{HashMap, HashSet};
use uuid::Uuid;

use crate::{Aggregate, Aggregator, GraphError, Meta, Metadata, MetadataFilter, Node, NodeStore};

/// An edge between two nodes in a graph.
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Edge {
    /// Instance metadata.
    pub metadata: Metadata,
    /// Source node ID.
    pub source: Uuid,
    /// Target node ID.
    pub target: Uuid,
}
impl Edge {
    /// Create a new edge struct with arguments.
    pub fn new(metadata: Option<Metadata>, source: Uuid, target: Uuid) -> Self {
        Edge {
            metadata: metadata.unwrap_or_default(),
            source: source.to_owned(),
            target: target.to_owned(),
        }
    }
    /// Create an edge between two newly created and also returned nodes.
    pub fn new_and_nodes(metadata: Option<Metadata>) -> (Self, Node, Node) {
        let source = Node::default();
        let target = Node::default();
        (
            Edge::new(metadata, source.id().to_owned(), target.id().to_owned()),
            source,
            target,
        )
    }
}

impl Meta for Edge {
    /// Get edge metadata.
    fn get_meta(&self) -> &Metadata {
        &self.metadata
    }
}

#[derive(Clone, Debug, Default, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct EdgeStoreData {
    edges: HashMap<Uuid, Edge>,
    forward: HashMap<Uuid, HashMap<Uuid, HashSet<Uuid>>>,
    reverse: HashMap<Uuid, HashMap<Uuid, HashSet<Uuid>>>,
    aggregate: Aggregate,
}
impl HasEdgeStore for EdgeStoreData {
    fn edge_store(&self) -> &EdgeStoreData {
        self
    }
    fn edge_store_mut(&mut self) -> &mut EdgeStoreData {
        self
    }
}
impl EdgeStore for EdgeStoreData {}

pub trait HasEdgeStore {
    /// Edge store reference.
    fn edge_store(&self) -> &EdgeStoreData;
    /// Mutable edge store reference.
    fn edge_store_mut(&mut self) -> &mut EdgeStoreData;
}

pub trait EdgeStore: HasEdgeStore {
    /// The number of edges in store.
    fn edges_len(&self) -> usize {
        self.edge_store().edges.len()
    }

    /// Whether the edge store is empty.
    fn is_edges_empty(&self) -> bool {
        self.edges_len() == 0
    }

    /// Add a single edge to this store.
    fn add_edge<N: NodeStore>(
        &mut self,
        edge: Edge,
        safe: bool,
        nodes: Option<&N>,
    ) -> Result<Option<Edge>> {
        if safe {
            match nodes {
                None => Err(anyhow!(GraphError::MissingNodeStore)),
                Some(nodes) => self.check_edge(&edge, nodes),
            }?;
        }

        let id = edge.id().to_owned();
        let replaced = self.del_edge(&id);
        let source = edge.source.to_owned();
        let target = edge.target.to_owned();

        let store = self.edge_store_mut();

        store.aggregate.add(edge.get_meta());
        store.edges.insert(id.to_owned(), edge);

        store
            .forward
            .entry(source.to_owned())
            .or_insert(HashMap::<Uuid, HashSet<Uuid>>::new())
            .entry(target.to_owned())
            .or_insert(HashSet::new())
            .insert(id.to_owned());

        store
            .reverse
            .entry(target)
            .or_insert(HashMap::<Uuid, HashSet<Uuid>>::new())
            .entry(source)
            .or_insert(HashSet::new())
            .insert(id);

        Ok(replaced)
    }

    /// Add multiple edges to this store at once.
    fn extend_edges<N: NodeStore>(
        &mut self,
        edges: Vec<Edge>,
        safe: bool,
        nodes: Option<&N>,
    ) -> Result<Vec<Edge>> {
        let mut replaced = vec![];
        let store = self.edge_store_mut();
        for edge in edges.into_iter() {
            if let Some(edge) = store.add_edge(edge, safe, nodes)? {
                replaced.push(edge);
            }
        }
        Ok(replaced)
    }

    /// Check an edge's source and target reference in a node store.
    fn check_edge<N: NodeStore>(&self, edge: &Edge, nodes: &N) -> Result<()> {
        if !nodes.has_node(&edge.source) {
            return Err(anyhow!(GraphError::NodeNotInStore(edge.source)));
        }
        if !nodes.has_node(&edge.target) {
            return Err(anyhow!(GraphError::NodeNotInStore(edge.target)));
        }
        Ok(())
    }

    /// Delete an edge from the store.
    fn del_edge(&mut self, id: &Uuid) -> Option<Edge> {
        let store = self.edge_store_mut();
        if let Some(edge) = store.edges.remove(id) {
            store.aggregate.subtract(edge.get_meta());
            let id = edge.id();
            let source = edge.source;
            let target = edge.target;

            if let Some(tgt_map) = store.forward.get_mut(&source) {
                if let Some(edge_ids) = tgt_map.get_mut(&target) {
                    edge_ids.remove(id);
                    if edge_ids.is_empty() {
                        tgt_map.remove(&target);
                    }
                }
                if tgt_map.is_empty() {
                    store.forward.remove(&source);
                }
            }

            if let Some(src_map) = store.reverse.get_mut(&target) {
                if let Some(edge_ids) = src_map.get_mut(&source) {
                    edge_ids.remove(id);
                    if edge_ids.is_empty() {
                        src_map.remove(&source);
                    }
                }
                if src_map.is_empty() {
                    store.reverse.remove(&target);
                }
            }

            return Some(edge);
        }
        None
    }

    /// Get an edge from the store.
    fn get_edge(&self, id: &Uuid) -> Option<&Edge> {
        self.edge_store().edges.get(id)
    }

    /// Get all edges in store.
    fn all_edges(&self) -> Vec<&Edge> {
        self.edge_store().edges.values().collect()
    }

    /// Get all edge IDs between a given source and target node.
    fn edge_ids_between(&self, source: &Uuid, target: &Uuid) -> HashSet<Uuid> {
        self.edge_store()
            .forward
            .get(source)
            .and_then(|tgt_map| tgt_map.get(target).cloned())
            .unwrap_or_default()
    }

    /// Get all edges between a given source and target node.
    fn edges_between(&self, source: &Uuid, target: &Uuid) -> Vec<&Edge> {
        self.edge_ids_between(source, target)
            .into_iter()
            .filter_map(|id| self.get_edge(&id))
            .collect()
    }

    /// Get all edge IDs between a given set of source and target nodes.
    fn edge_ids_between_all(
        &self,
        sources: &HashSet<Uuid>,
        targets: &HashSet<Uuid>,
    ) -> HashSet<Uuid> {
        let mut ids = HashSet::<Uuid>::new();
        for source in sources.iter() {
            for target in targets.iter() {
                ids.extend(self.edge_ids_between(source, target));
            }
        }
        ids
    }

    /// Get all edges between a given set of source and target nodes.
    fn edges_between_all(&self, sources: &HashSet<Uuid>, targets: &HashSet<Uuid>) -> Vec<&Edge> {
        self.edge_ids_between_all(sources, targets)
            .into_iter()
            .filter_map(|id| self.get_edge(&id))
            .collect()
    }

    /// Get the calculated aggregate of edge metadata between a source and target node
    /// with an optional edge filter.
    fn aggregate_between(
        &self,
        source: &Uuid,
        target: &Uuid,
        filter: Option<&MetadataFilter>,
    ) -> Aggregate {
        let mut agg = Aggregate::default();
        for edge in self.edges_between(source, target) {
            if let Some(filter) = filter {
                if !filter.satisfies(edge) {
                    continue;
                }
            }
            agg.add(edge.get_meta());
        }
        agg
    }

    /// Get the calculated aggregate of edge metadata between a set of source and target
    /// nodes with an optional edge filter.
    fn aggregate_between_all(
        &self,
        sources: &HashSet<Uuid>,
        targets: &HashSet<Uuid>,
        filter: Option<&MetadataFilter>,
    ) -> Aggregate {
        let mut agg = Aggregate::default();
        for source in sources.iter() {
            for target in targets.iter() {
                agg.extend(self.aggregate_between(source, target, filter));
            }
        }
        agg
    }

    /// Calculate an aggregate value between a source and target node, optional edge
    /// filter and optional field filter.
    fn aggregate_value_between(
        &self,
        source: &Uuid,
        target: &Uuid,
        aggregator: &Aggregator,
        filter: Option<&MetadataFilter>,
        fields: Option<&HashSet<String>>,
    ) -> f64 {
        let agg = self.aggregate_between(source, target, filter);
        agg.sum(aggregator, fields)
    }

    /// Get the aggregate map from source to target to aggregate for a given set of
    /// source and target nodes and an optional edge filter.
    fn aggregate_map(
        &self,
        sources: &HashSet<Uuid>,
        targets: &HashSet<Uuid>,
        filter: Option<&MetadataFilter>,
    ) -> HashMap<Uuid, HashMap<Uuid, Aggregate>> {
        let mut map = HashMap::new();
        for source in sources.iter() {
            for target in targets.iter() {
                let agg = self.aggregate_between(source, target, filter);
                map.entry(source.to_owned())
                    .or_insert(HashMap::<Uuid, Aggregate>::new())
                    .insert(target.to_owned(), agg);
            }
        }
        map
    }

    /// Get the aggregate matrix where row indices correspond to target nodes and column
    /// indices correspond to source nodes as given in their input vectors. It's
    /// optional to specify an edge filter.
    fn aggregate_matrix(
        &self,
        sources: &[&Uuid],
        targets: &[&Uuid],
        filter: Option<&MetadataFilter>,
    ) -> Vec<Vec<Aggregate>> {
        targets
            .iter()
            .map(|&target| {
                sources
                    .iter()
                    .map(|&source| self.aggregate_between(source, target, filter))
                    .collect()
            })
            .collect()
    }

    /// Get all outgoing edge IDs originating from this source node.
    fn outgoing_ids_from(&self, source: &Uuid) -> HashSet<Uuid> {
        let mut edge_ids = HashSet::new();
        if let Some(tgt_map) = self.edge_store().forward.get(source) {
            for bundle in tgt_map.values() {
                edge_ids.extend(bundle);
            }
        }
        edge_ids
    }

    /// Get all outgoing edges originating from this source node.
    fn outgoing_edges_from(&self, source: &Uuid) -> Vec<&Edge> {
        self.outgoing_ids_from(source)
            .iter()
            .filter_map(|id| self.get_edge(id))
            .collect()
    }

    /// Get all incoming edge IDs towards this target node.
    fn incoming_ids_to(&self, target: &Uuid) -> HashSet<Uuid> {
        let mut edge_ids = HashSet::new();
        if let Some(src_map) = self.edge_store().reverse.get(target) {
            for bundle in src_map.values() {
                edge_ids.extend(bundle);
            }
        }
        edge_ids
    }

    /// Get all incoming edges towards this target node.
    fn incoming_edges_to(&self, target: &Uuid) -> Vec<&Edge> {
        self.incoming_ids_to(target)
            .iter()
            .filter_map(|id| self.get_edge(id))
            .collect()
    }

    /// Get all source node IDs that are connected to this target node by an incoming
    /// edge.
    fn targets_of(&self, source: &Uuid) -> HashSet<Uuid> {
        let mut targets = HashSet::<Uuid>::new();
        if let Some(tgt_map) = self.edge_store().forward.get(source) {
            for bundle in tgt_map.values() {
                targets.extend(
                    bundle
                        .iter()
                        .filter_map(|edge_id| self.get_edge(edge_id).map(|edge| edge.target)),
                );
            }
        }
        targets
    }

    /// Get all source node IDs that are connected to this target node by an incoming
    /// edge.
    fn sources_to(&self, target: &Uuid) -> HashSet<Uuid> {
        if let Some(src_map) = self.edge_store().reverse.get(target) {
            let mut sources = HashSet::<Uuid>::new();
            for bundle in src_map.values() {
                sources.extend(
                    bundle
                        .iter()
                        .filter_map(|edge_id| self.get_edge(edge_id).map(|edge| edge.source)),
                );
            }
            return sources;
        }
        HashSet::new()
    }

    /// Check whether a node is connected to a set of other nodes. Optionally specify
    /// a limited set of edge IDs that are allowed for connections.
    fn is_connected_to(
        &self,
        node: &Uuid,
        others: &HashSet<Uuid>,
        edge_ids: Option<&HashSet<Uuid>>,
    ) -> bool {
        match edge_ids {
            None => {
                !self.targets_of(node).is_disjoint(others)
                    || !self.sources_to(node).is_disjoint(others)
            }
            Some(edge_ids) => {
                let connections =
                    self.edge_ids_between_all(&HashSet::from([node.to_owned()]), others);
                !connections.is_disjoint(edge_ids)
            }
        }
    }

    /// Calculate the adjacency matrix given a set of source and target nodes,
    /// aggregator, optional edge filter and optional field filter.
    fn adjacency_matrix(
        &self,
        sources: &Vec<&Uuid>,
        targets: &Vec<&Uuid>,
        aggregator: &Aggregator,
        filter: Option<&MetadataFilter>,
        fields: Option<&HashSet<String>>,
    ) -> DMatrix<f64> {
        let mut matrix = DMatrix::<f64>::zeros(targets.len(), sources.len());
        for (col, &source) in sources.iter().enumerate() {
            for (row, &target) in targets.iter().enumerate() {
                matrix[(row, col)] =
                    self.aggregate_value_between(source, target, aggregator, filter, fields)
            }
        }
        matrix
    }

    /// Get the aggregate.
    fn edge_aggregate(&self) -> &Aggregate {
        &self.edge_store().aggregate
    }
}