/*!

  Graph utils for the `graph` module.

*/

use crate::circuit::{Instantiable, Net};
use crate::error::Error;
#[cfg(feature = "graph")]
use crate::netlist::Connection;
use crate::netlist::{DrivenNet, InputPort, NetRef, Netlist};
#[cfg(feature = "graph")]
use petgraph::graph::DiGraph;
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, HashSet};

/// A common trait of analyses than can be performed on a netlist.
/// An analysis becomes stale when the netlist is modified.
pub trait Analysis<'a, I: Instantiable>
where
    Self: Sized + 'a,
{
    /// Construct the analysis to the current state of the netlist.
    fn build(netlist: &'a Netlist<I>) -> Result<Self, Error>;
}

/// A table that maps nets to the circuit nodes they drive
pub struct FanOutTable<'a, I: Instantiable> {
    /// A reference to the underlying netlist
    _netlist: &'a Netlist<I>,
    /// Maps a net to the list of nodes it drives
    net_fan_out: HashMap<Net, Vec<NetRef<I>>>,
    /// Maps a driven net to a list of inputs it drives
    dnet_fan_out: HashMap<DrivenNet<I>, Vec<InputPort<I>>>,
    /// Maps a node to the list of nodes it drives
    node_fan_out: HashMap<NetRef<I>, Vec<NetRef<I>>>,
    /// The number of references held by all the  data structures
    ref_count: HashMap<NetRef<I>, usize>,
    /// Contains nets which are outputs
    is_an_output: HashSet<Net>,
}

impl<I> FanOutTable<'_, I>
where
    I: Instantiable,
{
    /// Returns an iterator to the circuit nodes that use `net`.
    pub fn get_net_users(&self, net: &Net) -> impl Iterator<Item = NetRef<I>> {
        self.net_fan_out
            .get(net)
            .into_iter()
            .flat_map(|users| users.iter().cloned())
    }

    /// Returns an iterator to the circuit nodes that use `node`.
    pub fn get_node_users(&self, node: &NetRef<I>) -> impl Iterator<Item = NetRef<I>> {
        self.node_fan_out
            .get(node)
            .into_iter()
            .flat_map(|users| users.iter().cloned())
    }

    /// Returns an iterator to the uses of `net`.
    pub fn get_users(&self, net: &DrivenNet<I>) -> impl Iterator<Item = InputPort<I>> {
        self.dnet_fan_out
            .get(net)
            .into_iter()
            .flat_map(|users| users.iter().cloned())
    }

    /// Get the number of reference held by this table
    pub fn get_ref_count(&self, node: &NetRef<I>) -> usize {
        self.ref_count.get(node).copied().unwrap_or(0)
    }

    /// Returns `true` if the net has any used by any cells in the circuit
    /// This does incude nets that are only used as outputs.
    pub fn net_has_uses(&self, net: &Net) -> bool {
        (self.net_fan_out.contains_key(net) && !self.net_fan_out[net].is_empty())
            || self.is_an_output.contains(net)
    }

    /// Returns `true` if the net has any uses  in the circuit
    pub fn has_uses(&self, net: &DrivenNet<I>) -> bool {
        net.is_top_level_output()
            || (self.dnet_fan_out.contains_key(net) && !self.dnet_fan_out[net].is_empty())
    }
}

impl<'a, I> Analysis<'a, I> for FanOutTable<'a, I>
where
    I: Instantiable,
{
    fn build(netlist: &'a Netlist<I>) -> Result<Self, Error> {
        let mut net_fan_out: HashMap<Net, Vec<NetRef<I>>> = HashMap::new();
        let mut dnet_fan_out: HashMap<DrivenNet<I>, Vec<InputPort<I>>> = HashMap::new();
        let mut node_fan_out: HashMap<NetRef<I>, Vec<NetRef<I>>> = HashMap::new();
        let mut is_an_output: HashSet<Net> = HashSet::new();
        let mut ref_count: HashMap<NetRef<I>, usize> = HashMap::new();

        // We can only build the fanout table if netlist is mostly intact
        if let Err(e) = netlist.verify() {
            match e {
                Error::NoOutputs => (),
                _ => return Err(e),
            }
        }

        for c in netlist.connections() {
            let e = net_fan_out.entry(c.net()).or_default();
            e.push(c.target().unwrap());

            let e = dnet_fan_out.entry(c.src()).or_default();
            e.push(c.target());

            let e = node_fan_out.entry(c.src().unwrap()).or_default();
            e.push(c.target().unwrap());
        }

        for (o, n) in netlist.outputs() {
            is_an_output.insert(o.as_net().clone());
            is_an_output.insert(n);
        }

        for v in net_fan_out.values() {
            for nr in v {
                *ref_count.entry(nr.clone()).or_insert(1) += 1;
            }
        }

        for (k, v) in &dnet_fan_out {
            for nr in v {
                *ref_count.entry(nr.clone().unwrap()).or_insert(1) += 1;
            }
            *ref_count.entry(k.clone().unwrap()).or_insert(1) += 1;
        }

        for (k, v) in &node_fan_out {
            for nr in v {
                *ref_count.entry(nr.clone()).or_insert(1) += 1;
            }
            *ref_count.entry(k.clone()).or_insert(1) += 1;
        }

        Ok(FanOutTable {
            _netlist: netlist,
            net_fan_out,
            dnet_fan_out,
            node_fan_out,
            is_an_output,
            ref_count,
        })
    }
}

/// A simple example to analyze the logic levels of a netlist.
/// This analysis checks for cycles, but it doesn't check for registers.
/// Result of combinational depth analysis for a single net.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum CombDepthResult {
    /// Signal has no driver
    Undefined,
    /// Signal is along a cycle
    CombCycle,
    /// Integer logic level
    Depth(usize),
}

/// Computes the combinational depth of each net in a netlist.
///
/// Each net is classified as having a defined depth, being undefined,
/// or participating in a combinational cycle.
pub struct CombDepthInfo<'a, I: Instantiable> {
    _netlist: &'a Netlist<I>,
    /// The total distance from a sequential element
    results: HashMap<NetRef<I>, CombDepthResult>,
    /// The critical predecessor to the node
    critical_par: HashMap<NetRef<I>, InputPort<I>>,
    /// Critical endpoints to build paths from
    critical_ends: BinaryHeap<(Reverse<usize>, NetRef<I>)>,
    /// Max will be `None` if the entire circuit is part of a combinational cycle or has undriven elements
    max_depth: Option<usize>,
}

impl<I> CombDepthInfo<'_, I>
where
    I: Instantiable,
{
    /// Max number of critical endpoints to keep in the heap.
    const SIZE_HEAP: usize = 10;

    /// Returns the logic level of a node in the circuit.
    pub fn get_comb_depth(&self, node: &NetRef<I>) -> Option<CombDepthResult> {
        self.results.get(node).copied()
    }

    /// Returns the critical input port
    pub fn get_crit_input(&self, node: &NetRef<I>) -> Option<&InputPort<I>> {
        self.critical_par.get(node)
    }

    /// Returns the most critical endpoints in the circuit
    pub fn get_critical_points(&self) -> impl IntoIterator<Item = DrivenNet<I>> {
        let mut v = self.critical_ends.iter().collect::<Vec<_>>();
        v.sort_by_key(|(d, _)| *d);
        v.into_iter().flat_map(|(_, n)| n.outputs())
    }

    /// Builds the most critical path
    pub fn build_critical_path(&self) -> Option<Vec<DrivenNet<I>>> {
        let mut path = Vec::new();
        let mut current = self.get_critical_points().into_iter().next()?;
        while let Some(crit) = self.critical_par.get(&current.clone().unwrap()) {
            path.push(current.clone());
            current = self
                ._netlist
                .get_driver(current.unwrap(), crit.get_input_num())
                .unwrap();
        }
        path.push(current);
        Some(path)
    }

    /// Returns the maximum logic level of the circuit.
    pub fn get_max_depth(&self) -> Option<usize> {
        self.max_depth
    }
}

impl<'a, I> Analysis<'a, I> for CombDepthInfo<'a, I>
where
    I: Instantiable,
{
    fn build(netlist: &'a Netlist<I>) -> Result<Self, Error> {
        let mut results: HashMap<NetRef<I>, CombDepthResult> = HashMap::new();
        let mut critical_par: HashMap<NetRef<I>, InputPort<I>> = HashMap::new();
        let mut critical_ends: BinaryHeap<(_, NetRef<I>)> = BinaryHeap::new();
        let mut visiting: HashSet<NetRef<I>> = HashSet::new();
        let mut max_depth: Option<usize> = None;

        fn compute<I: Instantiable>(
            node: NetRef<I>,
            netlist: &Netlist<I>,
            results: &mut HashMap<NetRef<I>, CombDepthResult>,
            critical_par: &mut HashMap<NetRef<I>, InputPort<I>>,
            visiting: &mut HashSet<NetRef<I>>,
        ) -> CombDepthResult {
            // Memoized result
            if let Some(&r) = results.get(&node) {
                return r;
            }

            // Cycle detection
            if visiting.contains(&node) {
                for n in visiting.iter() {
                    results.insert(n.clone(), CombDepthResult::CombCycle);
                }
                return CombDepthResult::CombCycle;
            }

            // Input nodes and reg have depth 0
            if node.is_an_input() || node.get_instance_type().is_some_and(|inst| inst.is_seq()) {
                let r = CombDepthResult::Depth(0);
                results.insert(node.clone(), r);
                return r;
            }

            visiting.insert(node.clone());

            let mut max_depth = 0;
            let mut crit: Option<InputPort<I>> = None;
            let mut is_undefined = false;

            for i in 0..node.get_num_input_ports() {
                let driver = match netlist.get_driver(node.clone(), i) {
                    Some(d) => d.unwrap(),
                    None => {
                        is_undefined = true;
                        continue;
                    }
                };

                if let Some(inst) = driver.get_instance_type()
                    && inst.is_seq()
                {
                    continue;
                }

                match compute(driver, netlist, results, critical_par, visiting) {
                    CombDepthResult::Depth(d) => {
                        if d > max_depth {
                            max_depth = d;
                            crit = Some(node.get_input(i));
                        }
                    }
                    CombDepthResult::Undefined => {
                        is_undefined = true;
                    }
                    CombDepthResult::CombCycle => {
                        let r = CombDepthResult::CombCycle;
                        results.insert(node.clone(), r);
                        visiting.remove(&node);
                        return r;
                    }
                }
            }

            visiting.remove(&node);
            let r = if is_undefined {
                CombDepthResult::Undefined
            } else {
                if let Some(crit) = crit {
                    critical_par.insert(node.clone(), crit);
                }
                let d = max_depth + 1;
                CombDepthResult::Depth(d)
            };
            results.insert(node.clone(), r);
            r
        }

        for (driven, _) in netlist.outputs() {
            let node = driven.unwrap();
            let r = compute(
                node.clone(),
                netlist,
                &mut results,
                &mut critical_par,
                &mut visiting,
            );

            if let CombDepthResult::Depth(d) = r {
                critical_ends.push((Reverse(d), node));
                if critical_ends.len() > CombDepthInfo::<I>::SIZE_HEAP {
                    critical_ends.pop();
                }
                max_depth = Some(max_depth.map_or(d, |m| m.max(d)));
            }
        }

        for node in netlist.matches(|inst| inst.is_seq()) {
            compute(
                node.clone(),
                netlist,
                &mut results,
                &mut critical_par,
                &mut visiting,
            );
            for i in 0..node.get_num_input_ports() {
                if let Some(driver) = netlist.get_driver(node.clone(), i) {
                    if driver.get_instance_type().is_some_and(|inst| inst.is_seq()) {
                        continue;
                    }

                    let r = compute(
                        driver.clone().unwrap(),
                        netlist,
                        &mut results,
                        &mut critical_par,
                        &mut visiting,
                    );
                    if let CombDepthResult::Depth(d) = r {
                        critical_ends.push((Reverse(d), driver.unwrap()));
                        if critical_ends.len() > CombDepthInfo::<I>::SIZE_HEAP {
                            critical_ends.pop();
                        }
                        max_depth = Some(max_depth.map_or(d, |m| m.max(d)));
                    }
                }
            }
        }

        Ok(CombDepthInfo {
            _netlist: netlist,
            results,
            critical_par,
            critical_ends,
            max_depth,
        })
    }
}

/// An enum to provide pseudo-nodes for any misc user-programmable behavior.
#[cfg(feature = "graph")]
#[derive(Debug, Clone)]
pub enum Node<I: Instantiable, T: Clone + std::fmt::Debug + std::fmt::Display> {
    /// A 'real' circuit node
    NetRef(NetRef<I>),
    /// Any other user-programmable node
    Pseudo(T),
}

#[cfg(feature = "graph")]
impl<I, T> std::fmt::Display for Node<I, T>
where
    I: Instantiable,
    T: Clone + std::fmt::Debug + std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Node::NetRef(nr) => nr.fmt(f),
            Node::Pseudo(t) => std::fmt::Display::fmt(t, f),
        }
    }
}

/// An enum to provide pseudo-edges for any misc user-programmable behavior.
#[cfg(feature = "graph")]
#[derive(Debug, Clone)]
pub enum Edge<I: Instantiable, T: Clone + std::fmt::Debug + std::fmt::Display> {
    /// A 'real' circuit connection
    Connection(Connection<I>),
    /// Any other user-programmable node
    Pseudo(T),
}

#[cfg(feature = "graph")]
impl<I, T> std::fmt::Display for Edge<I, T>
where
    I: Instantiable,
    T: Clone + std::fmt::Debug + std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Edge::Connection(c) => c.fmt(f),
            Edge::Pseudo(t) => std::fmt::Display::fmt(t, f),
        }
    }
}

/// Returns a petgraph representation of the netlist as a directed multi-graph with type [DiGraph<Object, NetLabel>].
#[cfg(feature = "graph")]
pub struct MultiDiGraph<'a, I: Instantiable> {
    _netlist: &'a Netlist<I>,
    graph: DiGraph<Node<I, String>, Edge<I, Net>>,
}

#[cfg(feature = "graph")]
impl<I> MultiDiGraph<'_, I>
where
    I: Instantiable,
{
    /// Return a reference to the graph constructed by this analysis
    pub fn get_graph(&self) -> &DiGraph<Node<I, String>, Edge<I, Net>> {
        &self.graph
    }

    /// Iterates through a [greedy feedback arc set](https://doi.org/10.1016/0020-0190(93)90079-O) for the graph.
    pub fn greedy_feedback_arcs(&self) -> impl Iterator<Item = Connection<I>> {
        petgraph::algo::feedback_arc_set::greedy_feedback_arc_set(&self.graph)
            .map(|e| match e.weight() {
                Edge::Connection(c) => c,
                _ => unreachable!("Outputs should be sinks"),
            })
            .cloned()
    }

    /// Returns all the circuit nodes sorted into their strongly connected components.
    pub fn sccs(&self) -> Vec<Vec<NetRef<I>>> {
        let mut res = Vec::new();
        for scc in petgraph::algo::tarjan_scc(&self.graph) {
            let c: Vec<NetRef<I>> = scc
                .into_iter()
                .filter_map(|i| match &self.graph[i] {
                    Node::NetRef(nr) => Some(nr.clone()),
                    _ => None,
                })
                .collect();
            if !c.is_empty() {
                res.push(c);
            }
        }
        res
    }
}

#[cfg(feature = "graph")]
impl<'a, I> Analysis<'a, I> for MultiDiGraph<'a, I>
where
    I: Instantiable,
{
    fn build(netlist: &'a Netlist<I>) -> Result<Self, Error> {
        // If we verify, we can hash by name
        netlist.verify()?;
        let mut mapping = HashMap::new();
        let mut graph = DiGraph::new();

        for obj in netlist.objects() {
            let id = graph.add_node(Node::NetRef(obj.clone()));
            mapping.insert(obj.to_string(), id);
        }

        for connection in netlist.connections() {
            let source = connection.src().unwrap().get_obj().to_string();
            let target = connection.target().unwrap().get_obj().to_string();
            let s_id = mapping[&source];
            let t_id = mapping[&target];
            graph.add_edge(s_id, t_id, Edge::Connection(connection));
        }

        // Finally, add the output connections
        for (o, n) in netlist.outputs() {
            let s_id = mapping[&o.clone().unwrap().get_obj().to_string()];
            let t_id = graph.add_node(Node::Pseudo(format!("Output({n})")));
            graph.add_edge(s_id, t_id, Edge::Pseudo(o.as_net().clone()));
        }

        Ok(Self {
            _netlist: netlist,
            graph,
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{format_id, netlist::*};

    fn full_adder() -> Gate {
        Gate::new_logical_multi(
            "FA".into(),
            vec!["CIN".into(), "A".into(), "B".into()],
            vec!["S".into(), "COUT".into()],
        )
    }

    fn ripple_adder() -> GateNetlist {
        let netlist = Netlist::new("ripple_adder".to_string());
        let bitwidth = 4;

        // Add the the inputs
        let a = netlist.insert_input_logic_bus("a".to_string(), bitwidth);
        let b = netlist.insert_input_logic_bus("b".to_string(), bitwidth);
        let mut carry: DrivenNet<Gate> = netlist.insert_input("cin".into());

        for (i, (a, b)) in a.into_iter().zip(b).enumerate() {
            // Instantiate a full adder for each bit
            let fa = netlist
                .insert_gate(full_adder(), format_id!("fa_{i}"), &[carry, a, b])
                .unwrap();

            // Expose the sum
            fa.expose_net(&fa.get_net(0)).unwrap();

            carry = fa.find_output(&"COUT".into()).unwrap();

            if i == bitwidth - 1 {
                // Last full adder, expose the carry out
                fa.get_output(1).expose_with_name("cout".into()).unwrap();
            }
        }

        netlist.reclaim().unwrap()
    }

    #[test]
    fn fanout_table() {
        let netlist = ripple_adder();
        let analysis = FanOutTable::build(&netlist);
        assert!(analysis.is_ok());
        let analysis = analysis.unwrap();
        assert!(netlist.verify().is_ok());

        for item in netlist.objects().filter(|o| !o.is_an_input()) {
            // Sum bit has no users (it is a direct output)
            assert!(
                analysis
                    .get_net_users(&item.find_output(&"S".into()).unwrap().as_net())
                    .next()
                    .is_none(),
                "Sum bit should not have users"
            );

            assert!(
                item.get_instance_name().is_some(),
                "Item should have a name. Filtered inputs"
            );

            let net = item.find_output(&"COUT".into()).unwrap().as_net().clone();
            let mut cout_users = analysis.get_net_users(&net);
            if item.get_instance_name().unwrap().to_string() != "fa_3" {
                assert!(cout_users.next().is_some(), "Carry bit should have users");
            }

            assert!(
                cout_users.next().is_none(),
                "Carry bit should have 1 or 0 user"
            );
        }
    }
}
