// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
//
// SPDX-License-Identifier: AGPL-3.0-or-later

//! Construct timing graphs (delay graph, constraint graph) from a netlist.

use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;

use fnv::FnvHashMap;
use libreda_db::reference_access::*;
use libreda_db::traits::*;
use petgraph::visit::EdgeRef;
use smallvec::SmallVec;

use crate::traits::CellConstraintModel;
use crate::traits::ConstraintBase;
use crate::traits::{CellDelayModel, TimingBase};
use crate::StaError;
use crate::PATH_SEPARATOR;
use fnv::FnvHashSet;
use libreda_db::prelude as db;
use pargraph::id_rwlock::IdLockReadGuard;
use pargraph::BorrowDataCell;
use pargraph::DataCell;
use petgraph::prelude::NodeIndex;
use petgraph::stable_graph::StableDiGraph;

/// Data of a node in the timing graph.
/// Holds timing information such as actual and required arrival times but
/// also data used for concurrent graph operations.
#[derive(Debug)]
pub(crate) struct NodeData<N, D>
where
    N: NetlistIds,
    D: TimingBase + ConstraintBase,
{
    /// Type of the node. Also makes a link to the corresponding pin or pin-instance in the netlist.
    pub node_type: GraphNodeType<N>,
    /// Node data which can be accessed for read/write during parallel graph operations.
    sync_data: DataCell<SyncNodeData<D>>,
    /// The level of this node (minimal number of hops to a node without inputs)
    pub forward_level: AtomicU32,
    /// Count the number of unresolved input nodes.
    pub forward_dependencies: DependencyCounter,
    /// Count the number of unresolved output nodes + the current node.
    /// Note the node itself is considered a backward dependency because
    /// the back-propgation depends on the presence of the actual signal.
    pub backward_dependencies: DependencyCounter,
    /// ID of the last iteration in which this node was touched for forward propagation.
    /// Used to mark nodes for incremental updates.
    pub generation_forward: AtomicU32,
    /// ID of the last iteration in which this node was touched for backward propagation.
    /// Used to mark nodes for incremental updates.
    pub generation_backward: AtomicU32,
}

impl<N, D> NodeData<N, D>
where
    N: NetlistBase,
    D: TimingBase + ConstraintBase,
{
    pub fn new(node_type: GraphNodeType<N>) -> Self {
        Self {
            node_type,
            sync_data: Default::default(),
            forward_level: Default::default(),
            generation_forward: Default::default(),
            generation_backward: Default::default(),
            forward_dependencies: Default::default(),
            backward_dependencies: Default::default(),
        }
    }
}

impl<N, D> BorrowDataCell for NodeData<N, D>
where
    N: NetlistBase,
    D: TimingBase + ConstraintBase,
{
    type UserData = SyncNodeData<D>;

    fn borrow_data_cell(&self) -> &DataCell<Self::UserData> {
        &self.sync_data
    }
}

/// Atomic counter unresolved dependencies of graph nodes.
#[derive(Debug, Default)]
pub(crate) struct DependencyCounter {
    /// Number of unresolved dependencies.
    num_unresolved: AtomicU32,
}

impl DependencyCounter {
    /// Get number of unresolved dependencies.
    /// Loaded using relaxed memory ordering.
    pub fn num_unresolved(&self) -> u32 {
        self.num_unresolved.load(Ordering::Relaxed)
    }

    /// Set the number of unresolved dependencies.
    /// Uses relaxed memory ordering.
    pub fn set_num_unresolved(&self, num_unresolved: u32) {
        self.num_unresolved.store(num_unresolved, Ordering::Relaxed);
    }

    /// Increment the number of unresolved dependencies and return
    /// the result.
    pub fn increment_unresolved(&self) -> u32 {
        self.num_unresolved.fetch_add(1, Ordering::Relaxed) + 1
    }

    /// Decrement the number of unresolved dependencies and return
    /// the result.
    /// # Panics
    /// Panics on an integer underflow.
    pub fn decrement_unresolved(&self) -> u32 {
        let r = self.num_unresolved.fetch_sub(1, Ordering::Relaxed);
        assert!(r > 0, "integer underflow in number of dependencies");
        r - 1
    }
}

/// Node data which we can access during parallel graph operations.
/// Access to this data is synchronized by a lock in the `DataCell` struct.
#[derive(Debug, Clone)]
pub(crate) struct SyncNodeData<D>
where
    D: TimingBase + ConstraintBase,
{
    /// Mark this node as a primary input. Primary inputs
    /// have the actual signal set by the user. The actual signal of a primary input
    /// thus does not need to be computed.
    pub is_primary_input: bool,
    /// The signal at this location in the netlist.
    /// This is typically used to represent the arrival time.
    pub signal: Option<D::Signal>,
    /// The required signal which is necessary to meet constraints in the netlist.
    /// Typically this is encodes the required arrival time.
    pub required_signal: Option<D::RequiredSignal>,
}

impl<D> Default for SyncNodeData<D>
where
    D: TimingBase + ConstraintBase,
{
    fn default() -> Self {
        Self {
            signal: Default::default(),
            required_signal: Default::default(),
            is_primary_input: false,
        }
    }
}

/// Node in the timing graph.
/// Most nodes represent a pin in the netlist. There are also a few virtual nodes
/// (source nodes) which simplify the implementation of the delay propagation algorithms.
#[derive(Clone)]
pub(crate) enum GraphNodeType<N: NetlistIds> {
    /// A virtual node.
    /// Used as a start node in the delay propagation algorithm.
    ForwardPropagationSource,
    /// A virtual node which represents the starting node for backpropagating the constraints.
    BackwardPropagationSource,
    /// A node which represents a pin of a circuit component.
    /// Most of the nodes will be of this variant.
    Terminal(db::TerminalId<N>),
}

impl<N> std::fmt::Debug for GraphNodeType<N>
where
    N: NetlistIds,
    N::PinId: std::fmt::Debug,
    N::PinInstId: std::fmt::Debug,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            GraphNodeType::ForwardPropagationSource => write!(f, "source"),
            GraphNodeType::BackwardPropagationSource => write!(f, "sink"),
            GraphNodeType::Terminal(t) => write!(f, "Terminal({:?})", t),
        }
    }
}

#[derive(Debug)]
pub(crate) struct EdgeData<T: ConstraintBase> {
    pub edge_type: GraphEdgeType,
    pub sync_data: DataCell<SyncEdgeData<T>>,
}

impl<T> EdgeData<T>
where
    T: ConstraintBase,
{
    pub fn new(edge_type: GraphEdgeType) -> Self {
        Self {
            edge_type,
            sync_data: Default::default(),
        }
    }
}
impl<D> BorrowDataCell for EdgeData<D>
where
    D: TimingBase + ConstraintBase,
{
    type UserData = SyncEdgeData<D>;

    fn borrow_data_cell(&self) -> &DataCell<Self::UserData> {
        &self.sync_data
    }
}

/// Edge data which we can access during parallel graph operations.
/// Access to this data is synchronized by a lock in the `DataCell` struct.
#[derive(Debug, Clone)]
pub(crate) struct SyncEdgeData<D>
where
    D: TimingBase + ConstraintBase,
{
    pub delay: Option<D::Delay>,
    pub constraint: Option<D::Constraint>, // TODO: remove
}

impl<D> Default for SyncEdgeData<D>
where
    D: TimingBase + ConstraintBase,
{
    fn default() -> Self {
        Self {
            delay: Default::default(),
            constraint: Default::default(),
        }
    }
}

#[derive(Copy, Clone, PartialEq, Eq, Debug)]
pub(crate) enum GraphEdgeType {
    Delay,
    Constraint,
    Virtual,
}

impl GraphEdgeType {
    pub fn is_delay_arc(&self) -> bool {
        matches!(self, Self::Delay)
    }

    pub fn is_constraint_arc(&self) -> bool {
        matches!(self, Self::Constraint)
    }
}

#[derive(Debug)]
pub(crate) struct TimingGraph<N: NetlistIds, T: ConstraintBase> {
    /// Graph structure holding all delay and constraint arcs. Their value corresponds to the delay.
    pub(crate) arc_graph: StableDiGraph<NodeData<N, T>, EdgeData<T>>,
    /// The source node is used as a starting point for computing the actual arrival times.
    /// All primary inputs to the combinational logic are connected to this source node.
    pub(crate) aat_source_node: NodeIndex,
    /// The source node is used as a starting point for computing the required arrival times.
    /// All constrained pins are connected to this source node.
    pub(crate) rat_source_node: NodeIndex,
    /// Mapping from graph nodes to terminal IDs in the netlist.
    pub(crate) node2term: FnvHashMap<NodeIndex, db::TerminalId<N>>,
    /// Mapping from terminal IDs in the netlist to graph nodes.
    pub(crate) term2node: FnvHashMap<db::TerminalId<N>, NodeIndex>,
    /// List of nodes from which a timing change originates.
    /// They are used to find the forward-cone and backward-code which are
    /// subject to change during an incremental update.
    frontier: FnvHashSet<NodeIndex>,
}

impl<N: NetlistBase, T: ConstraintBase> TimingGraph<N, T> {
    /// Create empty timing graph.
    pub fn new() -> Self {
        let mut delay_arc_graph: StableDiGraph<_, _> = Default::default();
        let source_node =
            delay_arc_graph.add_node(NodeData::new(GraphNodeType::ForwardPropagationSource));
        let sink_node =
            delay_arc_graph.add_node(NodeData::new(GraphNodeType::BackwardPropagationSource));
        Self {
            arc_graph: delay_arc_graph,
            node2term: Default::default(),
            term2node: Default::default(),
            aat_source_node: source_node,
            rat_source_node: sink_node,
            frontier: Default::default(),
        }
    }

    /// Get read access to the node data.
    /// Returns `None` if the terminal does not exist in the timing graph.
    /// Panics, if the read-lock cannot be acquired. Use only while there's no timing propagation running!
    pub fn get_node_data(
        &self,
        pin: &db::TerminalId<N>,
    ) -> Option<IdLockReadGuard<SyncNodeData<T>>> {
        let graph_node = *self.term2node.get(pin)?;

        self.arc_graph
            .node_weight(graph_node)
            .expect("node has no weight")
            .borrow_data_cell()
            .try_read()
            .expect("failed to acquire read-lock")
            .into()
    }

    /// Get a reference to the delay graph.
    pub fn graph(&self) -> &StableDiGraph<NodeData<N, T>, EdgeData<T>> {
        &self.arc_graph
    }

    /// Take the set of frontier nodes and replace it by an empty set.
    pub fn take_frontier(&mut self) -> FnvHashSet<NodeIndex> {
        let frontier = std::mem::replace(&mut self.frontier, Default::default());
        frontier
    }

    /// Mark a node which needs to be recomputed in the next incremental update.
    pub fn add_to_frontier(&mut self, terminal: &db::TerminalId<N>) {
        let node = self
            .term2node
            .get(terminal)
            .expect("terminal does not exist in timing graph");
        self.frontier.insert(*node);
    }

    /// Assemble a directed graph which represents all delays arcs and constraints (without values yet).
    /// Net terminals become nodes in the graph.
    /// Cell-internal arcs are derived from the `cell_delay_arcs`.
    /// Inter-cell arcs (interconnects) are derived from the netlist of the top cell `top_ref`.
    pub fn build_from_netlist<C, D>(
        top_ref: &CellRef<N>,
        delay_model: &D,
        constraint_model: &C,
    ) -> Self
    where
        D: CellDelayModel<N, Delay = T::Delay>,
        C: CellConstraintModel<N, Constraint = T::Constraint>,
    {
        // Create empty timing graph.
        let mut timing_graph = Self::new();

        for p in top_ref.each_pin() {
            timing_graph.create_pin(p);
        }

        // Create nodes for each pin instance.
        for inst in top_ref.each_cell_instance() {
            timing_graph.create_cell_instance(&inst, delay_model, constraint_model);
        }

        timing_graph.create_delay_edges_inter_cell(top_ref);
        timing_graph.create_primary_input_delay_arcs(top_ref);

        // Add the root node to the frontier.
        // This makes sure that the first incremental signal propagation computes all signals.
        timing_graph.frontier.insert(timing_graph.aat_source_node);

        timing_graph
    }

    /// Create graph nodes for each pin instance.
    pub fn create_cell_instance<C, D>(
        &mut self,
        inst: &CellInstRef<N>,
        delay_model: &D,
        constraint_model: &C,
    ) where
        D: CellDelayModel<N, Delay = T::Delay>,
        C: CellConstraintModel<N, Constraint = T::Constraint>,
    {
        for p in inst.each_pin_instance() {
            let terminal = p.into_terminal();
            self.create_terminal(terminal.id());
        }

        self.create_delay_edges_intra_cell(inst, delay_model);
        self.create_constraint_edges(inst, constraint_model);
    }

    pub fn remove_cell_instance(&mut self, inst: &CellInstRef<N>) {
        for p in inst.each_pin_instance() {
            let terminal = p.into_terminal();
            self.remove_terminal(terminal.id());
        }
    }

    /// Remove all inter-cell arcs created by this net.
    pub fn remove_net(&mut self, net: &NetRef<N>) -> Result<(), StaError> {
        let terminals = net.each_terminal();
        // Distinguish between sources (drivers) and sinks.
        let (sources, sinks): (Vec<_>, Vec<_>) = terminals.into_iter().partition(|t| {
            // Is source?

            match t {
                db::TerminalRef::Pin(p) => p.direction().is_input(),
                db::TerminalRef::PinInst(p) => p.pin().direction().is_output(),
            }
        });

        if !sources.is_empty() {
            if sources.len() != 1 {
                log::warn!(
                    "Net must have 1 driver (has {}): {}",
                    sources.len(),
                    net.qname(PATH_SEPARATOR)
                );
                return Err(StaError::Other(format!(
                    "Net must have 1 driver (has {}): {}",
                    sources.len(),
                    net.qname(PATH_SEPARATOR)
                )));
            }

            let source = sources[0].id();
            let source_id = self.term2node[&source];

            for sink in sinks {
                let sink_id = self.term2node[&sink.id()];
                let edge = self
                    .arc_graph
                    .find_edge(source_id, sink_id)
                    .expect("delay arc is not present");
                self.arc_graph.remove_edge(edge);
            }
        }

        return Ok(());
    }

    /// Disconnect the terminal from the old net and connect it to the new net (if it is `Some`).
    pub fn connect_terminal(&mut self, t: db::TerminalRef<N>, net: Option<N::NetId>) {
        let is_source = match &t {
            db::TerminalRef::Pin(p) => p.direction().is_input(),
            db::TerminalRef::PinInst(p) => p.pin().direction().is_output(),
        };
        let node = *self
            .term2node
            .get(&t.id())
            .expect("terminal is not in timing graph");

        let dir = if is_source {
            petgraph::Direction::Outgoing
        } else {
            petgraph::Direction::Incoming
        };

        // Disconnect first.
        // Need to collect the edges because later we need mutable access to the graph.
        let edges: SmallVec<[_; 8]> = self
            .arc_graph
            .edges_directed(node, dir)
            // Sanity check.
            .inspect(|e| {
                assert_eq!(
                    e.weight().edge_type,
                    GraphEdgeType::Delay,
                    "should remove only interconnect delay arcs"
                )
            })
            .map(|e| (e.source(), e.target(), e.id()))
            .collect();

        let netlist = t.base();
        for (src, dest, edge_id) in edges {
            assert!(
                // Sanity check: make sure we don't accidentially remove an arc inside a cell but a interconnect arc.
                if is_source { src == node } else { dest == node },
                "this is not an interconnect edge"
            );
            self.arc_graph.remove_edge(edge_id);
            self.frontier.insert(src);
            self.frontier.insert(dest);
        }

        // Eventually connect to new net.
        if let Some(net) = net {
            let net = netlist.net_ref(&net);
            // Connect
            if is_source {
                for sink in net.each_sink() {
                    let sink_node = self.term2node[&sink.id()];
                    self.arc_graph.update_edge(
                        node,
                        sink_node,
                        EdgeData::new(GraphEdgeType::Delay),
                    );
                    self.frontier.insert(node);
                    self.frontier.remove(&sink_node); // Not necessary but might save some memory.
                }
            } else {
                for source in net.each_driver() {
                    let source_node = self.term2node[&source.id()];
                    self.arc_graph.update_edge(
                        source_node,
                        node,
                        EdgeData::new(GraphEdgeType::Delay),
                    );
                    self.frontier.insert(source_node);
                    self.frontier.remove(&node); // Not necessary but might save some memory.
                }
            };
        }
    }

    pub fn create_terminal(&mut self, terminal_id: db::TerminalId<N>) {
        debug_assert!(
            !self.term2node.contains_key(&terminal_id),
            "terminal is already present"
        );
        let node_id = self
            .arc_graph
            .add_node(NodeData::new(GraphNodeType::Terminal(terminal_id.clone())));
        self.term2node.insert(terminal_id.clone(), node_id);
        self.node2term.insert(node_id, terminal_id);
    }

    // Create grpah nodes for an external pin.
    fn create_pin(&mut self, pin: PinRef<N>) {
        let terminal = pin.into_terminal();
        self.create_terminal(terminal.id());
    }

    fn remove_terminal(&mut self, terminal: db::TerminalId<N>) {
        let node = self
            .term2node
            .remove(&terminal)
            .expect("pin does not exist in timing graph");
        self.arc_graph.remove_node(node);
        self.node2term.remove(&node).expect("node does not exist");
    }

    pub fn remove_pin(&mut self, pin: N::PinId) {
        let terminal = db::TerminalId::PinId(pin);
        self.remove_terminal(terminal)
    }

    pub fn remove_pin_instance(&mut self, pin: N::PinInstId) {
        let terminal = db::TerminalId::PinInstId(pin);
        self.remove_terminal(terminal)
    }

    fn create_delay_edges_inter_cell(&mut self, top_ref: &CellRef<N>) -> Result<(), StaError> {
        let netlist = top_ref.base();
        let top = top_ref.id();

        // Create inter-cell arcs.
        for net in netlist.each_internal_net(&top) {
            let terminals = netlist.each_terminal_of_net_vec(&net);
            // Distinguish between sources (drivers) and sinks.
            let (sources, sinks): (Vec<_>, Vec<_>) = terminals.into_iter().partition(|t| {
                // Is source?
                match t {
                    db::TerminalId::PinId(p) => netlist.pin_direction(p).is_input(),
                    db::TerminalId::PinInstId(p) => {
                        let pin = netlist.template_pin(p);
                        netlist.pin_direction(&pin).is_output()
                    }
                }
            });

            if !sources.is_empty() {
                if sources.len() != 1 {
                    log::warn!(
                        "Net must have 1 driver (has {}): {}",
                        sources.len(),
                        netlist.net_ref(&net).qname(PATH_SEPARATOR)
                    );
                    return Err(StaError::Other(format!(
                        "Net must have 1 driver (has {}): {}",
                        sources.len(),
                        netlist.net_ref(&net).qname(PATH_SEPARATOR)
                    )));
                }

                let source = &sources[0];
                let source_id = self.term2node[source];

                for sink in sinks {
                    let sink_id = self.term2node[&sink];
                    self.arc_graph.update_edge(
                        source_id,
                        sink_id,
                        EdgeData::new(GraphEdgeType::Delay),
                    );
                }
            }
        }

        Ok(())
    }

    fn pin_instance_to_graph_node(&self, pin_inst: &PinInstRef<N>) -> NodeIndex {
        let terminal = pin_inst.terminal_id();
        *self
            .term2node
            .get(&terminal)
            .expect("Pin node does not exist.")
    }

    fn create_delay_edges_intra_cell<D>(&mut self, inst: &CellInstRef<N>, delay_model: &D)
    where
        D: CellDelayModel<N, Delay = T::Delay>,
    {
        let netlist = inst.base();

        // Create intra-cell delay arcs.
        // Store the intra-cell delay edges.

        for delay_arc in delay_model.delay_arcs(netlist, &inst.template_id()) {
            // Convert terminals of arc to graph node indices.
            let [input_pin_node, output_pin_node] = [delay_arc.input_pin, delay_arc.output_pin]
                .map(|p| self.pin_instance_to_graph_node(&inst.pin_instance(&p)));

            self.arc_graph.update_edge(
                input_pin_node,
                output_pin_node,
                EdgeData::new(GraphEdgeType::Delay),
            );
        }
    }

    fn create_constraint_edges<C>(&mut self, inst: &CellInstRef<N>, constraint_model: &C)
    where
        C: CellConstraintModel<N, Constraint = T::Constraint>,
    {
        let netlist = inst.base();

        for constraint_arc in constraint_model.constraint_arcs(netlist, &inst.template_id()) {
            // Convert terminals of arc to graph node indices.
            let [related_pin_node, constrained_pin_node] =
                [constraint_arc.related_pin, constraint_arc.constrained_pin]
                    .map(|p| self.pin_instance_to_graph_node(&inst.pin_instance(&p)));

            self.arc_graph.update_edge(
                related_pin_node,
                constrained_pin_node,
                EdgeData::new(GraphEdgeType::Constraint),
            );

            //// Keep track of the constraint arc by adding an edge to the RAT source node.
            //self.arc_graph.update_edge(
            //    self.rat_source_node,
            //    constrained_pin_node,
            //    EdgeData::new(GraphEdgeType::Virtual),
            //);
        }
    }

    /// Create a directed edge from the AAT source node the primary input nodes.
    fn create_primary_input(&mut self, primary_input: db::TerminalId<N>) {
        let t_id = self.term2node[&primary_input];
        // TODO: Take input delays into account here.
        self.arc_graph.update_edge(
            self.aat_source_node,
            t_id,
            EdgeData::new(GraphEdgeType::Delay),
        );
    }

    /// Create directed edges from the AAT source node all primary input nodes.
    fn create_primary_input_delay_arcs(&mut self, top_cell: &CellRef<N>) {
        // Create directed edges from the AAT source node to all primary input nodes.
        for prim_input in get_primary_inputs(top_cell) {
            self.create_primary_input(prim_input.into_terminal().into());
        }
    }
}

/// Get input pins of a cell.
fn get_primary_inputs<'a, N: NetlistBase>(top_ref: &CellRef<'a, N>) -> Vec<PinRef<'a, N>> {
    top_ref
        .each_pin()
        .filter(|p| p.direction().is_input())
        .collect()
}