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

//! Propagate forward and backward cones from a set of frontier nodes.
//! The cones mark the regions where the timing must potentially be updated.

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

use libreda_db::traits::{NetlistBaseMT, NetlistIds};
use petgraph::{
    data::DataMap,
    stable_graph::NodeIndex,
    visit::{EdgeRef, GraphBase, IntoEdgesDirected},
};

use crate::{
    timing_graph::{EdgeData, GraphEdgeType, NodeData, TimingGraph},
    traits::ConstraintBase,
};

use pargraph::{
    executors::multi_thread::MultiThreadExecutor, local_view::LocalGraphView,
    worklists::biglock_fifo::FifoWorklist, worklists::WorklistPush, ReadonlyOperator,
};

/// Mark the fan-out cones of the supplied `frontier`-nodes with the current
/// `generation` number. Let the set of nodes in this fan-out cones be `C_fwd`.
/// Then mark also the nodes in the fanin-cones of the nodes in `C_fwd` with
/// the current `generation` number.
///
/// This is used as a preliminary step of incremental timing propagation.
/// Only the nodes from `C_fwd` will need to update their actual signals.
/// And only the nodes from `C_bwd` will need to update their required signals.
pub(crate) fn propagate_cones<Netlist, CellModel>(
    g: &TimingGraph<Netlist, CellModel>,
    frontier: impl Iterator<Item = NodeIndex>,
    generation: u32,
) where
    Netlist: NetlistBaseMT,
    CellModel: ConstraintBase,
{
    let operator = ConePropagationOp { generation };

    let executor = MultiThreadExecutor::new();

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

#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum PropagationDirection {
    Forward,
    Backward,
}

#[derive(Clone)]
pub(crate) struct ConePropagationOp {
    pub generation: u32,
}

/// 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
    }
}

impl<G, N, D> ReadonlyOperator<G> for ConePropagationOp
where
    N: NetlistIds,
    D: ConstraintBase,
    G: GraphBase
        + IntoEdgesDirected
        + DataMap<NodeWeight = NodeData<N, D>, EdgeWeight = EdgeData<D>>,
{
    type WorkItem = WorkItem<G::NodeId>;

    fn op(
        &self,
        work_item: Self::WorkItem,
        local_view: LocalGraphView<&G>,
        mut worklist: impl WorklistPush<Self::WorkItem>,
    ) {
        let local_view = &local_view;
        let active_node = local_view.active_node();

        // Find next neighbors depending on propagation direction.
        let next_neighbors = |direction: PropagationDirection| {
            let dir = match direction {
                PropagationDirection::Forward => petgraph::Direction::Outgoing,
                PropagationDirection::Backward => petgraph::Direction::Incoming,
            };

            local_view
                .edges_directed(active_node, dir)
                .map(move |e| match direction {
                    PropagationDirection::Forward => (e.target(), e.weight().edge_type),
                    PropagationDirection::Backward => (e.source(), e.weight().edge_type),
                })
        };

        let local_data = local_view.node_weight(active_node).unwrap();

        if work_item.dir == PropagationDirection::Forward {
            let is_forward_visited = local_data
                .generation_forward
                .swap(self.generation, Ordering::Relaxed)
                == self.generation;

            // Do forward propagation only, if the active node was
            // reached by forward propagation or if the active node
            // is an initial frontier node.
            if !is_forward_visited {
                // The active node must be forward-propagated before it can be backward-propagated.
                local_data.backward_dependencies.increment_unresolved();

                for (n, edge_type) in next_neighbors(PropagationDirection::Forward) {
                    let data = local_view.node_weight(n).unwrap();

                    // Keep track of the number of unresolved dependencies.
                    match edge_type {
                        GraphEdgeType::Delay | GraphEdgeType::Virtual => {
                            data.forward_dependencies.increment_unresolved();
                        }
                        GraphEdgeType::Constraint => {
                            data.backward_dependencies.increment_unresolved();
                        }
                    }

                    worklist.push(WorkItem::new_forward(n));
                }
            }
        }

        let is_backward_visited = local_data
            .generation_backward
            .swap(self.generation, Ordering::Relaxed)
            == self.generation;

        if is_backward_visited {
            // Nothing to do.
            return;
        }

        // Always do backward propagation.
        for (n, edge_type) in next_neighbors(PropagationDirection::Backward) {
            let data = local_view.node_weight(n).unwrap();

            // Keep track of the number of unresolved dependencies.
            match edge_type {
                GraphEdgeType::Delay | GraphEdgeType::Virtual => {
                    data.backward_dependencies.increment_unresolved();
                }
                GraphEdgeType::Constraint => {}
                GraphEdgeType::Virtual => {}
            }

            worklist.push(WorkItem::new_backward(n));
        }
    }
}