// Copyright (c) 2021 Weird Constructor <weirdconstructor@gmail.com>
// This file is a part of HexoDSP. Released under GPL-3.0-or-later.
// See README.md and COPYING for details.

use crate::dsp::NodeId;
use std::collections::HashMap;
use std::collections::HashSet;

pub const MAX_NODE_EDGES: usize = 64;
pub const UNUSED_NODE_EDGE: usize = 999999;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
struct Node {
    /// The [NodeId] of this node.
    node_id: NodeId,
    /// The output edges of this node.
    edges: [usize; MAX_NODE_EDGES],
    /// The first unused index in the `edges` array.
    unused_idx: usize,
}

impl Node {
    pub fn new() -> Self {
        Self { node_id: NodeId::Nop, edges: [UNUSED_NODE_EDGE; MAX_NODE_EDGES], unused_idx: 0 }
    }

    pub fn clear(&mut self) {
        self.node_id = NodeId::Nop;
        self.edges = [UNUSED_NODE_EDGE; MAX_NODE_EDGES];
        self.unused_idx = 0;
    }

    pub fn add_edge(&mut self, node_index: usize) {
        for ni in self.edges.iter().take(self.unused_idx) {
            if *ni == node_index {
                return;
            }
        }

        self.edges[self.unused_idx] = node_index;
        self.unused_idx += 1;
    }
}

#[derive(Debug, Clone)]
pub struct NodeGraphOrdering {
    node2idx: HashMap<NodeId, usize>,
    node_count: usize,
    nodes: Vec<Node>,

    in_degree: Vec<usize>,
}

impl NodeGraphOrdering {
    pub fn new() -> Self {
        Self {
            node2idx: HashMap::new(),
            node_count: 0,
            nodes: vec![Node::new(); 256],
            in_degree: vec![0; 256],
        }
    }

    pub fn clear(&mut self) {
        self.node2idx.clear();
        self.node_count = 0;
    }

    pub fn add_node(&mut self, node_id: NodeId) -> usize {
        if let Some(idx) = self.node2idx.get(&node_id) {
            *idx
        } else {
            let idx = self.node_count;
            self.node_count += 1;

            if self.nodes.len() < self.node_count {
                self.nodes.resize(self.node_count, Node::new());
                self.in_degree.resize(self.node_count, 0);
            }

            self.nodes[idx].clear();
            self.nodes[idx].node_id = node_id;
            self.node2idx.insert(node_id, idx);

            idx
        }
    }

    fn get_node(&self, node_id: NodeId) -> Option<&Node> {
        let idx = *self.node2idx.get(&node_id)?;
        Some(&self.nodes[idx])
    }

    fn get_node_mut(&mut self, node_id: NodeId) -> Option<&mut Node> {
        let idx = *self.node2idx.get(&node_id)?;
        Some(&mut self.nodes[idx])
    }

    pub fn add_edge(&mut self, from_node_id: NodeId, to_node_id: NodeId) {
        let to_idx = self.add_node(to_node_id);

        if let Some(from_node) = self.get_node_mut(from_node_id) {
            from_node.add_edge(to_idx);
        }
    }

    pub fn has_path(&self, from_node_id: NodeId, to_node_id: NodeId) -> Option<bool> {
        let mut visited_set: HashSet<NodeId> = HashSet::with_capacity(self.node_count);

        let mut node_stack = Vec::with_capacity(self.node_count);
        node_stack.push(from_node_id);

        while let Some(node_id) = node_stack.pop() {
            if visited_set.contains(&node_id) {
                return None;
            } else {
                visited_set.insert(node_id);
            }

            if node_id == to_node_id {
                return Some(true);
            }

            if let Some(node) = self.get_node(node_id) {
                for node_idx in node.edges.iter().take(node.unused_idx) {
                    node_stack.push(self.nodes[*node_idx].node_id);
                }
            }
        }

        return Some(false);
    }

    /// Run Kahn's Algorithm to find the node order for the directed
    /// graph. `out` will contain the order the nodes should be
    /// executed in. If `false` is returned, the graph contains cycles
    /// and no proper order can be computed. `out` will be cleared
    /// in this case.
    pub fn calculate_order(&mut self, out: &mut Vec<NodeId>) -> bool {
        let mut deq = std::collections::VecDeque::with_capacity(self.node_count);

        for indeg in self.in_degree.iter_mut() {
            *indeg = 0;
        }

        // Calculate the number of inputs for each node:
        for node in self.nodes.iter().take(self.node_count) {
            for out_node_idx in node.edges.iter().take(node.unused_idx) {
                self.in_degree[*out_node_idx] += 1;
            }
        }

        // Find the "input" nodes by looking at the nodes without any input. These
        // are the first that must be executed!
        for idx in 0..self.node_count {
            if self.in_degree[idx] == 0 {
                deq.push_back(idx);
            }
        }

        // Track visit count for cycle detection:
        let mut visited_count = 0;

        while let Some(node_idx) = deq.pop_front() {
            visited_count += 1;

            let node = &self.nodes[node_idx];

            out.push(node.node_id);

            // Push the nodes we output to, to the end of the dequeue:
            for neigh_node_idx in node.edges.iter().take(node.unused_idx) {
                // by reducing the input count of the to be visited node:
                self.in_degree[*neigh_node_idx] -= 1;

                // we know when that node has been fully supplied with all it's inputs,
                // so we push it to the end of the dequeue to be executed then:
                if self.in_degree[*neigh_node_idx] == 0 {
                    deq.push_back(*neigh_node_idx);
                }
            }
        }

        if visited_count != self.node_count {
            out.clear();
            false
        } else {
            true
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn check_ngraph_dfs_1() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));
        ng.add_node(NodeId::Sin(0));

        ng.add_edge(NodeId::Sin(2), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(0), NodeId::Sin(1));

        assert!(ng.has_path(NodeId::Sin(2), NodeId::Sin(1)).unwrap());
        assert!(ng.has_path(NodeId::Sin(2), NodeId::Sin(0)).unwrap());
        assert!(ng.has_path(NodeId::Sin(0), NodeId::Sin(1)).unwrap());
        assert!(!ng.has_path(NodeId::Sin(1), NodeId::Sin(0)).unwrap());
        assert!(!ng.has_path(NodeId::Sin(0), NodeId::Sin(2)).unwrap());
        assert!(!ng.has_path(NodeId::Amp(0), NodeId::Out(2)).unwrap());
    }

    #[test]
    fn check_ngraph_order_1() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));

        ng.add_edge(NodeId::Sin(2), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(0), NodeId::Sin(1));

        let mut out = vec![];
        assert!(ng.calculate_order(&mut out));
        assert_eq!(out[..], [NodeId::Sin(2), NodeId::Sin(0), NodeId::Sin(1)]);
    }

    #[test]
    fn check_ngraph_order_2() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));
        ng.add_node(NodeId::Out(0));
        ng.add_node(NodeId::Amp(0));
        ng.add_node(NodeId::Amp(1));

        ng.add_edge(NodeId::Sin(2), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(0), NodeId::Sin(1));

        let mut out = vec![];
        assert!(ng.calculate_order(&mut out));
        assert_eq!(
            out[..],
            [
                NodeId::Sin(2),
                NodeId::Out(0),
                NodeId::Amp(0),
                NodeId::Amp(1),
                NodeId::Sin(0),
                NodeId::Sin(1)
            ]
        );
    }

    #[test]
    fn check_ngraph_order_3() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));
        ng.add_node(NodeId::Out(0));
        ng.add_node(NodeId::Amp(0));
        ng.add_node(NodeId::Amp(1));

        /*
            amp0 =>         sin0
            sin2 =>         sin0 => sin1 => out0
                 => amp1 => sin0
        */

        ng.add_edge(NodeId::Sin(2), NodeId::Sin(0));
        ng.add_edge(NodeId::Amp(0), NodeId::Sin(0));
        ng.add_edge(NodeId::Amp(1), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(2), NodeId::Amp(1));

        ng.add_edge(NodeId::Sin(0), NodeId::Sin(1));
        ng.add_edge(NodeId::Sin(1), NodeId::Out(0));

        let mut out = vec![];
        assert!(ng.calculate_order(&mut out));
        assert_eq!(
            out[..],
            [
                NodeId::Sin(2),
                NodeId::Amp(0),
                NodeId::Amp(1),
                NodeId::Sin(0),
                NodeId::Sin(1),
                NodeId::Out(0),
            ]
        );
    }

    #[test]
    fn check_ngraph_order_4() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));
        ng.add_node(NodeId::Out(0));
        ng.add_node(NodeId::Amp(0));
        ng.add_node(NodeId::Amp(1));

        /*
            amp1 => amp0 => sin0
            sin2 => sin1 => out0
        */

        ng.add_edge(NodeId::Amp(1), NodeId::Amp(0));
        ng.add_edge(NodeId::Amp(0), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(2), NodeId::Sin(1));
        ng.add_edge(NodeId::Sin(1), NodeId::Out(0));

        let mut out = vec![];
        assert!(ng.calculate_order(&mut out));
        assert_eq!(
            out[..],
            [
                NodeId::Sin(2),
                NodeId::Amp(1),
                NodeId::Sin(1),
                NodeId::Amp(0),
                NodeId::Out(0),
                NodeId::Sin(0),
            ]
        );
    }

    #[test]
    fn check_ngraph_dfs_cycle_2() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));

        ng.add_edge(NodeId::Sin(2), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(0), NodeId::Sin(1));
        ng.add_edge(NodeId::Sin(0), NodeId::Sin(2));

        assert!(ng.has_path(NodeId::Sin(2), NodeId::Sin(1)).is_none());

        let mut out = vec![];
        assert!(!ng.calculate_order(&mut out));
    }

    #[test]
    fn check_ngraph_clear() {
        let mut ng = NodeGraphOrdering::new();
        ng.add_node(NodeId::Sin(2));
        ng.add_node(NodeId::Sin(1));
        ng.add_node(NodeId::Sin(0));

        ng.add_edge(NodeId::Sin(2), NodeId::Sin(0));
        ng.add_edge(NodeId::Sin(0), NodeId::Sin(1));

        assert!(ng.has_path(NodeId::Sin(2), NodeId::Sin(1)).unwrap());

        ng.clear();

        assert!(!ng.has_path(NodeId::Sin(2), NodeId::Sin(1)).unwrap());
    }
}