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

//! Graph levelization based on an operator formulation.
//! Compute the level of all graph nodes in a directed acyclic graph.
//! The level of a node is the minimum distance to a node without incoming edges with each edge having a distance of `1`.

use num_traits::Zero;
use smallvec::*;

use std::{borrow::Borrow, marker::PhantomData, sync::atomic::Ordering};

use libreda_db::{prelude as db, reference_access::NetlistReferenceAccess};
use petgraph::{
    data::DataMap,
    visit::{Data, EdgeRef, GraphBase, IntoEdgesDirected},
};

use crate::{
    cone_propagation::{self, ConePropagationOp, PropagationDirection},
    timing_graph::*,
    traits::*,
};

use pargraph::{
    executors::multi_thread::MultiThreadExecutor,
    local_view::{FullConflictDetection, LocalGraphView, SyncReadGuard},
    worklists::{biglock_fifo::FifoWorklist, WorklistPush},
    BorrowDataCell, DataConflictErr, LabellingOperator,
};

/// Incrementally propagate actual and required signals in the full timing graph.
pub(crate) fn propagate_signals_incremental<Netlist, CellModel, ICModel, ICLoadModel>(
    netlist: &Netlist,
    g: &TimingGraph<Netlist, CellModel>,
    cell_model: &CellModel,
    interconnect_delay_model: &ICModel,
    interconnect_load_model: &ICLoadModel,
    frontier: impl Iterator<Item = petgraph::stable_graph::NodeIndex> + Clone,
    generation: u32,
) where
    Netlist: db::NetlistBaseMT,
    CellModel: ConstraintBase + CellDelayModel<Netlist> + CellConstraintModel<Netlist> + Sync,

    ICModel: InterconnectDelayModel<Netlist, Signal = CellModel::Signal> + Sync,
    ICLoadModel: InterconnectLoadModel<Netlist, Load = CellModel::Load> + Sync,
{
    let executor = MultiThreadExecutor::new();
    debug_assert!(
        frontier.clone().all(|n| {
            g.arc_graph
                .node_weight(n)
                .unwrap()
                .generation_forward
                .load(Ordering::Relaxed)
                < generation
        }),
        "`generation` must monotonically increase with each call"
    );

    // Mark cones which must be updated.
    {
        let operator = ConePropagationOp { generation };

        let worklist = FifoWorklist::new_with_local_queues(
            frontier
                .clone()
                .map(cone_propagation::WorkItem::new_forward)
                .collect(),
        );
        executor.run_readonly(worklist, operator, &g.arc_graph);
    }

    // Find root nodes of the cones. I.e. nodes which have all dependencies resolved.
    // They are a subset of the frontier.
    let initial_nodes = frontier
        .filter(|n| {
            let data = g.arc_graph.node_weight(*n).unwrap();

            data.forward_dependencies.num_unresolved() == 0
        })
        .map(WorkItem::new_forward);

    // Propagate signals in the cones.
    let propagation_op = PropagationOp::new(
        netlist,
        cell_model,
        interconnect_delay_model,
        interconnect_load_model,
        generation,
    );

    let worklist = FifoWorklist::new_global_queue(initial_nodes.collect());
    executor.run_node_labelling(worklist, &propagation_op, &g.arc_graph);
}

/// Propagate actual signals, constraints and required signals on the full (!) netlist.
/// Requires that all (!) primary inputs of the delay graph are present in the initial worklist. Otherwise the algorithm will miss
/// parts of the graph.
#[derive(Clone)]
struct PropagationOp<'a, G, N, D, IC, ICL> {
    /// Eventually need the netlist to do the mapping from terminals to cells.
    /// TODO: remove
    netlist: &'a N,
    /// The cell model provides the timing behaviour of cells (delays, constraints).
    cell_model: &'a D,
    /// Model for wire delays.
    interconnect_delay_model: &'a IC,
    /// Provides the loads imposed by nets attached to the outputs of logic gates.
    interconnect_load_model: &'a ICL,
    /// The number of the current iteration.
    /// Used to detect nodes which have not yet been touched in the current iteration.
    generation: u32,
    _graph_type: PhantomData<G>,
    _node_data_type: PhantomData<N>,
}

impl<'a, G, N, C, IC, ICL> PropagationOp<'a, G, N, C, IC, ICL>
where
    N: db::NetlistBase,
    C: DelayBase + ConstraintBase + CellDelayModel<N> + CellConstraintModel<N>,
    IC: InterconnectDelayModel<N, Signal = C::Signal>,
    ICL: InterconnectLoadModel<N, Load = C::Load>,
    G: GraphBase
        + Data<NodeWeight = NodeData<N, C>, EdgeWeight = EdgeData<C>>
        + IntoEdgesDirected
        + DataMap,
    G::NodeId: std::fmt::Debug,
{
    fn new(
        netlist: &'a N,
        delay_model: &'a C,
        interconnect_delay_model: &'a IC,
        interconnect_load_model: &'a ICL,
        generation: u32,
    ) -> Self {
        Self {
            generation,
            netlist,
            cell_model: delay_model,
            _node_data_type: Default::default(),
            _graph_type: Default::default(),
            interconnect_delay_model,
            interconnect_load_model,
        }
    }

    /// Propagate the signal along a delay arc.
    /// The graph edge represents the delay arc. The delay arc might go from a cell output
    /// to a cell input (inter-cell or wiring delay) or the other way from a cell input to a cell output (intra-cell delay).
    fn evaluate_delay_arc(
        &self,
        netlist: &N,
        local_view: &LocalGraphView<&G, FullConflictDetection>,
        input_signal: &C::Signal,
        output_load: &C::Load,
        graph_edge: G::EdgeRef,
        other_inputs: &impl Fn(&N::PinId) -> Option<C::LogicValue>,
    ) -> Option<C::Signal> {
        let source = graph_edge.source();
        let target = graph_edge.target();

        let get_terminal = |node_id| {
            let node = &local_view
                .node_weight(node_id)
                .expect("node has no weight")
                .node_type;
            match node {
                GraphNodeType::Terminal(t) => Some(netlist.terminal_ref(t)),
                _ => None,
            }
        };

        let source_node = get_terminal(source)?;
        let target_node = get_terminal(target)?;

        match (&source_node, &target_node) {
            (db::TerminalRef::PinInst(s), db::TerminalRef::PinInst(t)) => {
                // Intra-cell delay
                let input_pin = s.pin().id();
                let output_pin = t.pin().id();

                let is_intra_cell =
                    s.pin().direction().is_input() && t.pin().direction().is_output();

                if is_intra_cell {
                    assert_eq!(
                        s.cell_instance(),
                        t.cell_instance(),
                        "pins must belong to the same cell instance"
                    );

                    self.cell_model.cell_output(
                        self.netlist,
                        &input_pin,
                        input_signal,
                        &output_pin,
                        output_load,
                        other_inputs,
                    )
                } else {
                    // Inter-cell
                    let output_load = Zero::zero(); // TODO
                    self.interconnect_delay_model.interconnect_output(
                        netlist,
                        &source_node.id(),
                        input_signal,
                        &target_node.id(),
                        &output_load,
                    )
                }
            }
            _ => {
                // Interconnect delay.
                let output_load = Zero::zero(); // TODO
                self.interconnect_delay_model.interconnect_output(
                    netlist,
                    &source_node.id(),
                    input_signal,
                    &target_node.id(),
                    &output_load,
                )
            }
        }
    }

    fn evaluate_constraint_arc(
        &self,
        local_view: &LocalGraphView<&G, FullConflictDetection>,
        local_data: &SyncNodeData<C>,
        constraint_edge: G::EdgeRef,
        other_inputs: &impl Fn(&N::PinId) -> Option<C::Signal>,
        output_loads: &impl Fn(&N::PinId) -> Option<C::Load>,
    ) -> Option<C::RequiredSignal> {
        let netlist = self.netlist;

        let constrained_node = constraint_edge.target();
        let related_node = constraint_edge.source();

        let constrained_node_weight = local_view.node_weight(constrained_node).unwrap();
        let related_node_weight = local_view.node_weight(related_node).unwrap();

        let constrained_node_data = local_data;

        let related_node_data = related_node_weight
            .borrow_data_cell()
            .try_read()
            .expect("data of constrained node is locked");

        // Translate a node ID to a netlist terminal.
        let get_terminal = |node_type: &_| match node_type {
            GraphNodeType::Terminal(t) => Some(netlist.terminal_ref(t)),
            _ => None,
        };

        let constrained_terminal = get_terminal(&constrained_node_weight.node_type)?;
        let related_terminal = get_terminal(&related_node_weight.node_type)?;

        match (&related_terminal, &constrained_terminal) {
            (db::TerminalRef::PinInst(r), db::TerminalRef::PinInst(c)) => {
                // Intra-cell constraint arc.
                let related_pin = r.pin().id();
                let constrained_pin = c.pin().id();

                let is_intra_cell =
                    r.pin().direction().is_input() && c.pin().direction().is_input();

                if is_intra_cell {
                    assert_eq!(
                        r.cell_instance(),
                        c.cell_instance(),
                        "pins must belong to the same cell instance"
                    );

                    let constrained_pin_signal = constrained_node_data
                        .signal
                        .as_ref()
                        .expect("actual arrival time at constrained pin is not known yet");
                    let related_pin_signal = related_node_data
                        .signal
                        .as_ref()
                        .expect("actual arrival time at related pin is not known yet");

                    self.cell_model.get_required_input(
                        netlist,
                        &c.pin().id(),
                        constrained_pin_signal,
                        &r.pin().id(),
                        related_pin_signal,
                        other_inputs,
                        output_loads,
                    )
                } else {
                    // Inter-cell constraint.
                    let output_load = C::Load::zero(); // TODO
                    todo!()
                }
            }
            _ => {
                // Interconnect delay.
                todo!("interconnect constraint");
            }
        }
    }

    /// Mark this dependency as resolved.
    /// Create new activities from output nodes which have all dependencies resolved.
    fn mark_forward_dependency_resolved(
        &self,
        local_view: &LocalGraphView<&G, FullConflictDetection>,
        push_to_worklist: &mut impl FnMut(WorkItem<G::NodeId>),
    ) {
        let output_edges =
            local_view.edges_directed(local_view.active_node(), petgraph::Direction::Outgoing);

        for e in output_edges {
            let output_node = e.target();
            let edge_type = e.weight().edge_type;

            let unsync_data = local_view
                .node_weight(output_node)
                .expect("output node has no weight");

            match edge_type {
                GraphEdgeType::Delay => {
                    let num_unresolved_dependencies =
                        unsync_data.forward_dependencies.decrement_unresolved();

                    if num_unresolved_dependencies == 0 {
                        // This thread was the last dependency. The output node can now be processed.
                        push_to_worklist(WorkItem::new_forward(output_node));
                    }
                }
                GraphEdgeType::Constraint => {}
                GraphEdgeType::Virtual => {}
            }

            match edge_type {
                GraphEdgeType::Virtual => {}
                GraphEdgeType::Delay => {}
                GraphEdgeType::Constraint => {
                    let num_unresolved_dependencies =
                        unsync_data.backward_dependencies.decrement_unresolved();

                    if num_unresolved_dependencies == 0 {
                        // This thread was the last dependency. The output node can now be processed.
                        push_to_worklist(WorkItem::new_backward(output_node));
                    }
                }
            }
        }

        // Mark the active node as resolved for back propagation.
        {
            let num_unresolved = local_view
                .node_weight(local_view.active_node())
                .unwrap()
                .backward_dependencies
                .decrement_unresolved();

            if num_unresolved == 0 {
                push_to_worklist(WorkItem::new_backward(local_view.active_node()))
            }
        }
    }

    /// Mark this dependency as resolved for backward propagation.
    /// Create new activities from output nodes which have all dependencies resolved.
    fn mark_backward_dependency_resolved(
        &self,
        local_view: &LocalGraphView<&G, FullConflictDetection>,
        push_to_worklist: &mut impl FnMut(WorkItem<G::NodeId>),
    ) {
        let input_edges = local_view
            .edges_directed(local_view.active_node(), petgraph::Direction::Incoming)
            .filter(|e| e.weight().edge_type.is_delay_arc());

        for e in input_edges {
            let input_node = e.source();

            let unsync_data = local_view
                .node_weight(input_node)
                .expect("input node has no weight");

            match unsync_data.node_type {
                GraphNodeType::Terminal(_) => {}
                _ => continue, // Skip virtual nodes.
            }

            let num_unresolved_dependencies =
                unsync_data.backward_dependencies.decrement_unresolved();

            if num_unresolved_dependencies == 0 {
                // This thread was the last dependency. The output node can now be processed.
                push_to_worklist(WorkItem::new_backward(input_node));
            }
        }
    }

    fn propagate_forward(
        &self,
        local_view: &LocalGraphView<&G, FullConflictDetection>,
        local_data: &mut SyncNodeData<C>,
        mut push_to_worklist: impl FnMut(WorkItem<G::NodeId>),
    ) -> Result<(), DataConflictErr<G::NodeId, G::EdgeId>> {
        let active_node = local_view.active_node();
        let unsync_data = local_view.node_weight(active_node).unwrap();

        debug_assert_eq!(
            unsync_data.generation_forward.load(Ordering::Relaxed),
            self.generation,
            "node is not marked for backward propagation"
        );

        // Compute the actual signal.
        if !local_data.is_primary_input {
            let input_edges = local_view
                .edges_directed(active_node, petgraph::Direction::Incoming)
                .filter(|e| e.weight().edge_type.is_delay_arc());

            // Prefetch the node data of the inputs and acquire the read locks.
            let inputs_with_data: Result<
                SmallVec<[(G::EdgeRef, SyncReadGuard<G::NodeWeight>); 8]>,
                _,
            > = input_edges
                .map(|e| {
                    local_view
                        .try_node_weight(e.source())
                        .expect("node has no weight")
                        .map(|data| (e, data))
                })
                .collect();

            let inputs_with_data = inputs_with_data?; // Abort on eventual errors. Should not happen.

            // Get the load connected to the node.
            let output_load = match &unsync_data.node_type {
                GraphNodeType::Terminal(t) => self
                    .interconnect_load_model
                    .interconnect_load(self.netlist, t),
                _ => Zero::zero(),
            };

            let other_inputs = |pin: &N::PinId| -> Option<C::LogicValue> {
                inputs_with_data.iter().find_map(|(edge_ref, data)| {
                    let src = edge_ref.source();
                    let node_data = local_view.node_weight(src).unwrap();
                    let found = match &node_data.node_type {
                        GraphNodeType::Terminal(t) => match t {
                            db::TerminalId::PinId(p) => pin == p,
                            db::TerminalId::PinInstId(p_inst) => {
                                &self.netlist.template_pin(p_inst) == pin
                            }
                        },
                        _ => false,
                    };

                    found
                        .then(|| data.signal.as_ref())
                        .flatten()
                        .map(|s| s.logic_value())
                })
            };

            local_data.signal = inputs_with_data
                .iter()
                // Propagate the signal along all incoming edges.
                .filter_map(|(input_edge, input_node_data)| {
                    input_node_data.signal.as_ref().and_then(|input_signal| {
                        let output_signal = self.evaluate_delay_arc(
                            self.netlist,
                            local_view,
                            input_signal,
                            &output_load,
                            *input_edge,
                            &other_inputs,
                        );

                        // Store the delay in the edge data.
                        let mut w = input_edge
                            .weight()
                            .borrow_data_cell()
                            .try_write(0)
                            // This algorithm should not have more than one concurrent write access to an edge.
                            .expect("failed to acquire write lock");
                        w.delay = output_signal
                            .as_ref()
                            .map(|o| self.cell_model.get_delay(input_signal, o));

                        output_signal
                    })
                })
                // Join the incoming signals into one signal (this is typically a max/min function of the arrival time).
                .reduce(|s1, s2| self.cell_model.summarize_delays(&s1, &s2));
        } else {
            // The node is a primary input.
            assert!(
                local_data.signal.is_some(),
                "Primary inputs must have their actual signal specified."
            );
        }

        self.mark_forward_dependency_resolved(local_view, &mut push_to_worklist);

        Ok(())
    }

    /// Compute the required signals such that constraints are met.
    /// Signal requirements of a node are imposed by outgoing delay arcs and incoming constraint arcs.
    fn propagate_backward(
        &self,
        local_view: &LocalGraphView<&G, FullConflictDetection>,
        local_data: &mut SyncNodeData<C>,
        mut push_to_worklist: impl FnMut(WorkItem<G::NodeId>),
    ) -> Result<(), DataConflictErr<G::NodeId, G::EdgeId>> {
        let active_node = local_view.active_node();
        let unsync_data = local_view.node_weight(active_node).unwrap();

        debug_assert_eq!(
            unsync_data.generation_backward.load(Ordering::Relaxed),
            self.generation,
            "node is not marked for backward propagation"
        );

        let output_delay_arcs = local_view
            .edges_directed(active_node, petgraph::Direction::Outgoing)
            .filter(|e| e.weight().edge_type.is_delay_arc());

        let input_constraint_arcs = local_view
            .edges_directed(active_node, petgraph::Direction::Incoming)
            .filter(|e| e.weight().edge_type.is_constraint_arc());

        // Prefetch the node data of the dependencies and acquire the read locks.
        let outputs_with_data: Result<
            SmallVec<[(G::EdgeRef, SyncReadGuard<G::NodeWeight>); 8]>,
            _,
        > = output_delay_arcs
            .map(|e| {
                local_view
                    .try_node_weight(e.target())
                    .expect("node has no weight")
                    .map(|data| (e, data))
            })
            .collect();
        let outputs_with_data = outputs_with_data?; // Abort on eventual errors.

        // Prefetch the node data of the dependencies and acquire the read locks.
        let constraint_inputs_with_data: Result<
            SmallVec<[(G::EdgeRef, SyncReadGuard<G::NodeWeight>); 8]>,
            _,
        > = input_constraint_arcs
            .map(|e| {
                local_view
                    .try_node_weight(e.source())
                    .expect("node has no weight")
                    .map(|data| (e, data))
            })
            .collect();
        assert!(constraint_inputs_with_data.is_ok());
        let constraint_inputs_with_data = constraint_inputs_with_data?; // Abort on eventual errors. Should not happen.

        // Resolve delay arcs.
        let required_signals_from_delays = outputs_with_data
            .into_iter()
            // Process only delay arcs which impose a constraint.
            .filter(|(_, output_data)| output_data.required_signal.is_some())
            .map(|(delay_edge, output_data)| {
                // Fetch the stored delay from the edge data.
                let edge_weight = delay_edge
                    .weight()
                    .borrow_data_cell()
                    .try_read()
                    .expect("Edge data should not be locked anymore.");

                let actual_delay = edge_weight
                    .delay
                    .as_ref()
                    .expect("Delay should already be computed.");

                let required_output = output_data.required_signal.as_ref().unwrap(); // Required signals which are `None` should be filtered out.

                let actual_input = local_data
                    .signal
                    .as_ref()
                    .expect("Actual signal is not computed yet.");

                self.cell_model
                    .solve_delay_constraint(actual_delay, required_output, actual_input)
            });

        // Resolve constraint arcs.
        let required_signals_from_constraints =
            constraint_inputs_with_data.into_iter().filter_map(
                |(constraint_edge, constraint_node_data)| {
                    self.evaluate_constraint_arc(
                        local_view,
                        local_data,
                        constraint_edge,
                        &|_pin| None,
                        &|pin| {
                            Some(self.interconnect_load_model.interconnect_load(
                                self.netlist,
                                &db::TerminalId::PinId(pin.clone()),
                            ))
                        },
                    )
                },
            );

        // Summarize all requirements into a single one.
        let required_signal = required_signals_from_delays
            .chain(required_signals_from_constraints)
            .reduce(|r1, r2| self.cell_model.summarize_constraints(&r1, &r2));

        // Store the required signal in the node's data.
        local_data.required_signal = required_signal;

        self.mark_backward_dependency_resolved(local_view, &mut push_to_worklist);
        Ok(())
    }
}

impl<'a, G, N, D, IC, ICL> LabellingOperator<G> for &PropagationOp<'a, G, N, D, IC, ICL>
where
    N: db::NetlistBase,
    D: DelayBase + ConstraintBase + CellDelayModel<N> + CellConstraintModel<N>,
    IC: InterconnectDelayModel<N, Signal = D::Signal>,
    ICL: InterconnectLoadModel<N, Load = D::Load>,
    G: GraphBase
        + Data<NodeWeight = NodeData<N, D>, EdgeWeight = EdgeData<D>>
        + IntoEdgesDirected
        + DataMap,
    G::NodeId: std::fmt::Debug,
{
    type NodeWeight = SyncNodeData<D>;
    type WorkItem = WorkItem<G::NodeId>;

    fn op(
        &self,
        work_item: Self::WorkItem,
        local_view: LocalGraphView<&G, FullConflictDetection>,
        local_data: &mut Self::NodeWeight,
        mut worklist: impl WorklistPush<Self::WorkItem>,
    ) -> Result<(), DataConflictErr<G::NodeId, G::EdgeId>> {
        match work_item.dir {
            PropagationDirection::Forward => {
                self.propagate_forward(&local_view, local_data, |n| worklist.push(n))?;
            }
            PropagationDirection::Backward => {
                self.propagate_backward(&local_view, local_data, |n| worklist.push(n))?;
            }
        };

        Ok(())
    }
}

/// Work item associated with a propagation direction.
#[derive(Clone, Copy, Debug)]
pub(crate) struct WorkItem<N> {
    node: N,
    dir: PropagationDirection,
}

impl<N> WorkItem<N> {
    pub fn new_forward(n: N) -> Self {
        Self {
            node: n,
            dir: PropagationDirection::Forward,
        }
    }
    fn new_backward(n: N) -> Self {
        Self {
            node: n,
            dir: PropagationDirection::Backward,
        }
    }
}

impl<N> Borrow<N> for WorkItem<N> {
    fn borrow(&self) -> &N {
        &self.node
    }
}