// Copyright (c) 2022-2023, NVIDIA CORPORATION. All rights reserved.
//
// NVIDIA CORPORATION and its licensors retain all intellectual property
// and proprietary rights in and to this software, related documentation
// and any modifications thereto. Any use, reproduction, disclosure or
// distribution of this software and related documentation without an express
// license agreement from NVIDIA CORPORATION is strictly prohibited.
//

//! @file Traversal.h
//!
//! @brief Defines graph traversal algorithms.
//! See @rstref{Traversing a Graph <ef_graph_traversal_guide>} and
//! @rstref{Graph Traversal In-Depth <ef_graph_traversal_advanced>}
//! for more information how to utilize graph traversals.
#pragma once

#include <concurrentqueue/include_concurrentqueue.h>
#include <omni/graph/exec/unstable/AtomicBackoff.h>
#include <omni/graph/exec/unstable/EnumBitops.h>

#include <atomic>
#include <functional>
#include <memory>
#include <queue>
#include <type_traits>
#include <vector>

namespace omni
{
namespace graph
{
namespace exec
{
namespace unstable
{

//! Implementation details.
namespace detail
{

//! Information returned to traversal callback indicating visit order
enum class VisitOrder
{
    eUnknown = 0, //!< Visit order is not specified
    eFirst = 1 << 0, //!< This is the first visit to a node during the traversal
    eNext = 1 << 1, //!< This is a next visit to a node, i.e. not a first one and not the last one
    eLast = 1 << 2, //!< This is the last visit to a node during the traversal
    eCycle = 1 << 3 //!< In case of cycles, when traversal enters node at the last visit, more visit can happen with
                    //!< cycle order
};

} // namespace detail

//! Enable bitwise operation on VisitOrder
template <>
struct EnumBitops<detail::VisitOrder> : EnumBitops<>::allow_bitops
{
};

namespace detail
{

//! Traversal information stored per node
struct NodeData
{
    std::atomic<std::size_t> visitCount{ 0 }; //!< How many times this node was visited

    //! Copy constructor
    //!
    //! Compiler will not generate by default copy constructor (nor assignment operator) due to atomic member.
    //! We are adding it explicitly here because the usage pattern is not allowing copy construction
    //! in concurrent execution.
    NodeData(const NodeData& src) : visitCount(src.visitCount.load())
    {
    }
    //! Assignment operator
    //!
    //! Compiler will not generate by default copy constructor (nor assignment operator) due to atomic member.
    //! We are adding it explicitly here because the usage pattern is not allowing assignment
    //! in concurrent execution.
    NodeData& operator=(const NodeData& rhs)
    {
        visitCount.store(rhs.visitCount.load());
        return *this;
    }
    // and because we added above, we still have to add default constructor which can't be added implicitly
    //! Default constructor
    NodeData()
    {
    }
};

// ef-docs visit-first-begin
//! Traversal strategy that enters the node when it was first time discovered
struct VisitFirst
{
    //! Call to traverse the graph with a strategy to visit only when first time discovered
    template <typename Node, typename NodeData>
    static VisitOrder tryVisit(Node* node, NodeData& nodeData)
    {
        auto lastVisit = nodeData.visitCount++; // read+increment only once. other threads can be doing the same.
        return (lastVisit == 0) ? VisitOrder::eFirst : VisitOrder::eUnknown;
    }
};
// ef-docs visit-first-end

// ef-docs visit-last-begin
//! Traversal strategy that enters the node when entire upstream was already visited and this is the last
//! opportunity to enter the node.
//!
//! In case of cycles, this algorithm is relying on knowledge of number of parents that are causing cycles.
struct VisitLast
{
    //! Call to traverse the graph with a strategy to visit only when no more visits are possible
    template <typename Node, typename NodeData>
    static VisitOrder tryVisit(Node& node, NodeData& nodeData)
    {
        auto requiredCount = node->getParents().size() - node->getCycleParentCount();
        auto currentVisit = ++nodeData.visitCount; // increment+read only once. other threads can be doing the same.
        if (requiredCount == 0 && currentVisit == 1)
        {
            return VisitOrder::eLast;
        }
        else if (currentVisit == requiredCount)
        {
            return VisitOrder::eLast;
        }

        return VisitOrder::eUnknown;
    }
};
// ef-docs visit-last-end

// ef-docs visit-all-begin
//! Traversal strategy that allows discovering all the edges in the graph. Traversal continuation is controlled by user
//! code.
struct VisitAll
{
    //! Call to traverse the graph with a strategy to visit all edges of the graph
    template <typename Node, typename NodeData>
    static VisitOrder tryVisit(Node& node, NodeData& nodeData)
    {
        auto parentCount = node->getParents().size();
        auto requiredCount = parentCount - node->getCycleParentCount();

        auto currentVisit = ++nodeData.visitCount; // increment+read only once. other threads can be doing the same.
        if (requiredCount == 0 && currentVisit == 1)
        {
            return (VisitOrder::eFirst | VisitOrder::eLast);
        }

        VisitOrder ret = VisitOrder::eUnknown;
        if (currentVisit > requiredCount)
        {
            ret = VisitOrder::eCycle;
        }
        else if (currentVisit == requiredCount)
        {
            ret = (currentVisit == 1) ? (VisitOrder::eFirst | VisitOrder::eLast) : VisitOrder::eLast;
        }
        else if (currentVisit == 1)
        {
            ret = VisitOrder::eFirst;
        }
        else
        {
            ret = VisitOrder::eNext;
        }

        return ret;
    }
};
// ef-docs visit-all-end

#ifndef DOXYGEN_BUILD

struct FlowDFS
{
};
struct FlowBFS
{
};

struct SerialQueue
{
};
struct ConcurrentQueue
{
};

template <typename GraphNode, typename Flow, typename Queue, bool ConstNode, typename Enable = void>
struct TraversalBase
{
};

template <typename GraphNode, typename Flow, typename Queue, bool ConstNode>
struct TraversalBase<GraphNode, Flow, Queue, ConstNode, std::enable_if_t<std::is_same<Flow, FlowDFS>::value>>
{
    void incrementInfo()
    {
    }
    void decrementInfo()
    {
    }
};

//! Base traversal class when using BFS traversal order and doesn't require thread safe queue
template <typename GraphNode, typename Flow, typename Queue, bool ConstNode>
struct TraversalBase<GraphNode,
                     Flow,
                     Queue,
                     ConstNode,
                     std::enable_if_t<std::is_same<Flow, FlowBFS>::value && std::is_same<Queue, SerialQueue>::value>>
{
    using Node = typename std::conditional_t<ConstNode, const GraphNode, GraphNode>;
    using NodeQueue = std::queue<Node*>;

    NodeQueue m_queue;

    void push(Node* node)
    {
        m_queue.push(&node);
    }
    bool tryPop(Node*& node)
    {
        if (m_queue.empty())
        {
            return false;
        }

        node = m_queue.front();
        m_queue.pop();
        return true;
    }

    void incrementInfo()
    {
    }
    void decrementInfo()
    {
    }
    bool hasInfo() const
    {
        return false;
    }
};

//! Base traversal class when using BFS traversal order and DOES require thread safe queue
template <typename GraphNode, typename Flow, typename Queue, bool ConstNode>
struct TraversalBase<GraphNode,
                     Flow,
                     Queue,
                     ConstNode,
                     std::enable_if_t<std::is_same<Flow, FlowBFS>::value && std::is_same<Queue, ConcurrentQueue>::value>>
{
    using Node = typename std::conditional_t<ConstNode, const GraphNode, GraphNode>;
    using NodeQueue = moodycamel::ConcurrentQueue<Node*>;

    NodeQueue m_queue;
    std::atomic<std::size_t> m_infoCount{ 0 };

    TraversalBase() noexcept
    {
    }
    TraversalBase(TraversalBase&& src) noexcept : m_queue(std::move(src.m_queue)), m_infoCount(src.m_infoCount.load())
    {
    }

    TraversalBase(TraversalBase& src) = delete;

    void push(Node* node)
    {
        m_queue.enqueue(node);
    }
    bool tryPop(Node*& node)
    {
        return m_queue.try_dequeue(node);
    }

    void incrementInfo()
    {
        m_infoCount++;
    }
    void decrementInfo()
    {
        m_infoCount--;
    }
    bool hasInfo() const
    {
        return m_infoCount.load() > 0;
    }
};

struct NoUserData
{
};

template <typename GraphNode, typename NodeUserData, typename Enable = void>
struct UserDataBase
{
    explicit UserDataBase(std::size_t size) noexcept
    {
    }
    UserDataBase(UserDataBase&& src) noexcept
    {
    }
};

//! User defined data class to be available for each traversed node
template <typename GraphNode, typename NodeUserData>
struct UserDataBase<GraphNode, NodeUserData, std::enable_if_t<!std::is_same<NodeUserData, NoUserData>::value>>
{
    static_assert(std::is_trivially_copyable<NodeUserData>::value, "User data needs to be trivially copyable");

    explicit UserDataBase(std::size_t size) noexcept : m_userData(size)
    {
    }

    UserDataBase(UserDataBase&& src) noexcept : m_userData(std::move(src.m_userData))
    {
    }

    NodeUserData& userData(GraphNode* node)
    {
        return m_userData[node->getIndexInTopology()];
    }

    using NodeUserDataArray = std::vector<NodeUserData>;
    NodeUserDataArray m_userData;
};

using QueueType = ConcurrentQueue; // or SerialQueue (but serial queue will fail with some multithreaded unit tests)

//! Traversal class
//!
//! @tparam GraphNode       Node typename
//! @tparam Graph           Graph typename
//! @tparam Strategy        Traversal visit strategy (first, last, all) typename
//! @tparam Flow            Visit flow (DFS or BFS) typename
//! @tparam NodeUserData    Custom user data typename allocated for each node
//! @tparam ConstNode       Is this a const traversal
template <typename GraphNode, typename Graph, typename Strategy, typename Flow, typename NodeUserData, bool ConstNode>
class Traversal : private TraversalBase<GraphNode, Flow, QueueType, ConstNode>,
                  private UserDataBase<GraphNode, NodeUserData>
{
public:
    using Node = typename std::conditional_t<ConstNode, const GraphNode, GraphNode>;
    using Base = TraversalBase<GraphNode, Flow, QueueType, ConstNode>;
    using BaseUserData = UserDataBase<GraphNode, NodeUserData>;

    struct Info
    {
        Traversal& traversal;
        VisitOrder order;

        Info(Traversal& t, VisitOrder o) noexcept : traversal(t), order(o)
        {
            traversal.Base::incrementInfo();
        }
        Info(const Info& src) noexcept : traversal(src.traversal), order(src.order)
        {
            traversal.Base::incrementInfo();
        }
        Info(Info&& src) noexcept : traversal(src.traversal), order(src.order)
        {
            traversal.Base::incrementInfo();
        }
        ~Info()
        {
            traversal.Base::decrementInfo();
        }

        Info& operator=(const Info& rhs) = delete;
        Info& operator=(Info&& rhs) = delete;

        bool isFirstVisit() const
        {
            return (order & VisitOrder::eFirst) == VisitOrder::eFirst;
        }
        bool isLastVisit() const
        {
            return (order & VisitOrder::eLast) == VisitOrder::eLast;
        }

        void continueVisit(GraphNode* node)
        {
            return traversal.continueVisit(node);
        }

        NodeUserData& userData(GraphNode* node)
        {
            return traversal.BaseUserData::userData(node);
        }
    };

    using CallbackType = void(Info, GraphNode*, GraphNode*);
    using CallbackFn = std::function<CallbackType>;
    using NodeDataArray = std::vector<NodeData>;

    explicit Traversal(Graph* g, CallbackFn call) noexcept
        : BaseUserData(g->getNodeCount()), m_callback(call), m_data(g->getNodeCount())
    {
    }

    Traversal() = delete;
    Traversal(const Traversal& src) = delete;
    Traversal(Traversal&& src) = delete;
    Traversal& operator=(const Traversal& rhs) = delete;
    Traversal& operator=(Traversal&& rhs) = delete;

    void continueVisit(GraphNode* prev)
    {
        continueImpl(Flow(), prev);
    }

    void startVisit(GraphNode* node)
    {
        startImpl(Flow(), node);
    }

    void markVisited(GraphNode* node)
    {
        nodeData(node).visitCount = node->getParents().size();
    }

private:
    // see if our traversal policy allows us to visit this node
    void tryVisit(GraphNode* prev, GraphNode* current)
    {
        VisitOrder visitOrder = Strategy::tryVisit(current, nodeData(current));
        if (visitOrder > VisitOrder::eUnknown)
        {
            m_callback({ *this, visitOrder }, prev, current);
        }
    }

    void tryContinue(GraphNode* current)
    {
        for (auto child : current->getChildren())
        {
            tryVisit(current, child);
        }
    }

    void continueImpl(FlowDFS, Node* prev)
    {
        tryContinue(prev);
    }

    void continueImpl(FlowBFS, Node* prev)
    {
        Base::push(prev);
    }

    void startImpl(FlowDFS, Node* node)
    {
        markVisited(node);
        tryContinue(node);
    }

    void startImpl(FlowBFS, Node* node)
    {
        markVisited(node);
        Base::push(node);

        AtomicBackoff backoff;
        Node* stackNode = nullptr;
        while (true)
        {
            if (Base::tryPop(stackNode))
            {
                tryContinue(stackNode);
                backoff.reset();
                continue;
            }

            if (Base::hasInfo())
            {
                backoff.pause();
            }
            else
            {
                break;
            }
        }
    }

    NodeData& nodeData(GraphNode* node)
    {
        return m_data[node->getIndexInTopology()];
    }

    CallbackFn m_callback;
    NodeDataArray m_data;
};

//! Main traversal template for DFS algorithms
template <typename GraphNode, typename Graph, typename Strategy, typename NodeUserData, bool ConstNode>
void run_traversal_dfs(typename std::conditional_t<ConstNode, const GraphNode, GraphNode>* node,
                       typename Traversal<GraphNode, Graph, Strategy, FlowDFS, NodeUserData, ConstNode>::CallbackFn call)
{
    Traversal<GraphNode, Graph, Strategy, FlowDFS, NodeUserData, ConstNode> traversal(node->getTopology(), call);
    traversal.startVisit(node);
}

//! Main traversal template for BFS algorithms
template <typename GraphNode, typename Graph, typename Strategy, typename NodeUserData, bool ConstNode>
void run_traversal_bfs(typename std::conditional_t<ConstNode, const GraphNode, GraphNode>* node,
                       typename Traversal<GraphNode, Graph, Strategy, FlowBFS, NodeUserData, ConstNode>::CallbackFn call)
{
    Traversal<GraphNode, Graph, Strategy, FlowBFS, NodeUserData, ConstNode> traversal(node->getTopology(), call);
    traversal.startVisit(node);
}

//! Main traversal template for DFS algorithms. Traversal is allocated on the heap and returned.
template <typename GraphNode, typename Graph, typename Strategy, typename NodeUserData, bool ConstNode>
auto alloc_and_run_traversal_dfs(
    typename std::conditional_t<ConstNode, const GraphNode, GraphNode>* node,
    typename Traversal<GraphNode, Graph, Strategy, FlowDFS, NodeUserData, ConstNode>::CallbackFn call)
{
    using TraversalType = Traversal<GraphNode, Graph, Strategy, FlowDFS, NodeUserData, ConstNode>;
    std::unique_ptr<TraversalType> traversal = std::make_unique<TraversalType>(node->getTopology(), call);
    traversal->startVisit(node);
    return traversal;
}

//! Main traversal template for BFS algorithms. Traversal is allocated on the heap and returned.
template <typename GraphNode, typename Graph, typename Strategy, typename NodeUserData, bool ConstNode>
auto alloc_and_run_traversal_bfs(
    typename std::conditional_t<ConstNode, const GraphNode, GraphNode>* node,
    typename Traversal<GraphNode, Graph, Strategy, FlowBFS, NodeUserData, ConstNode>::CallbackFn call)
{
    using TraversalType = Traversal<GraphNode, Graph, Strategy, FlowBFS, NodeUserData, ConstNode>;
    std::unique_ptr<TraversalType> traversal = std::make_unique<TraversalType>(node->getTopology(), call);
    traversal->startVisit(node);
    return traversal;
}

#endif // DOXYGEN_BUILD

} // detail
} // namespace unstable
} // namespace exec
} // namespace graph
} // namespace omni

#include <omni/graph/exec/unstable/IGraph.h>
#include <omni/graph/exec/unstable/INode.h>
#include <omni/graph/exec/unstable/ITopology.h>

namespace omni
{
namespace graph
{
namespace exec
{
namespace unstable
{

//! @copydoc omni::graph::exec::unstable::detail::VisitFirst
using VisitFirst = detail::VisitFirst;
//! @copydoc omni::graph::exec::unstable::detail::VisitLast
using VisitLast = detail::VisitLast;
//! @copydoc omni::graph::exec::unstable::detail::VisitAll
using VisitAll = detail::VisitAll;

using detail::FlowBFS;
using detail::FlowDFS;
using detail::NoUserData;
using detail::Traversal;
using detail::VisitOrder;

// ef-docs traversal-methods-begin
//! Deep-first-search traversal
template <typename Strategy, typename NodeUserData = NoUserData>
void traversal_dfs(INode* node,
                   typename Traversal<INode, ITopology, Strategy, FlowDFS, NodeUserData, false>::CallbackFn call)
{
    detail::run_traversal_dfs<INode, ITopology, Strategy, NodeUserData, false>(node, call);
}

//! Breadth-first-search traversal
template <typename Strategy, typename NodeUserData = NoUserData>
void traversal_bfs(INode* node,
                   typename Traversal<INode, ITopology, Strategy, FlowBFS, NodeUserData, false>::CallbackFn call)
{
    detail::run_traversal_bfs<INode, ITopology, Strategy, NodeUserData, false>(node, call);
}

//! Deep-first-search traversal. Traversal allocated on heap and returned to extend lifetime
//! to the end of all concurrent tasks.
template <typename Strategy, typename NodeUserData = NoUserData>
auto concurrent_traversal_dfs(INode* node,
                              typename Traversal<INode, ITopology, Strategy, FlowDFS, NodeUserData, false>::CallbackFn call)
{
    return detail::alloc_and_run_traversal_dfs<INode, ITopology, Strategy, NodeUserData, false>(node, call);
}

//! Breadth-first-search traversal. Traversal allocated on heap and returned to extend lifetime
//! to the end of all concurrent tasks.
template <typename Strategy, typename NodeUserData = NoUserData>
auto concurrent_traversal_bfs(INode* node,
                              typename Traversal<INode, ITopology, Strategy, FlowBFS, NodeUserData, false>::CallbackFn call)
{
    return detail::alloc_and_run_traversal_bfs<INode, ITopology, Strategy, NodeUserData, false>(node, call);
}
// ef-docs traversal-methods-end

} // namespace unstable
} // namespace exec
} // namespace graph
} // namespace omni
