// 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 GraphUtils.h
//!
//! @brief Defines utilities for graph visualization.
#pragma once

#include <carb/Format.h>

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

#include <ostream>
#include <string>
#include <unordered_map>
#include <vector>

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

//! Debugging utility to write out the graph topology in a graphviz format
//!
//! @param  inGraph Graph to dump
//! @param  out     Output stream to receive the produced graphviz text output
inline void writeFlattenedAsGraphviz(omni::core::ObjectParam<IGraph> inGraph, std::ostream& out);

namespace detail
{

#ifndef DOXYGEN_BUILD

constexpr const char* colorScheme()
{
    return "paired10";
}

constexpr unsigned colorSchemeSize()
{
    return 10;
}

struct GraphState
{
    std::unordered_map<size_t, unsigned> colorMapping;
    unsigned nextColorIndex{ 0 };

    unsigned getColor(size_t hash)
    {
        /*static const std::vector<const char*> gColors = { "black","aqua","aquamarine","bisque",
            "blue","blueviolet","brown","burlywood","cadetblue","chartreuse","chocolate","coral","cornflowerblue",
            "crimson","cyan","darkblue","darkcyan","darkgoldenrod","darkgray","darkgreen","darkgrey","darkkhaki",
            "darkmagenta","darkolivegreen","darkorange","darkorchid","darkred","darksalmon","darkseagreen","darkslateblue",
            "darkslategray","darkslategrey","darkturquoise","darkviolet","deeppink","deepskyblue","dimgray","dimgrey","dodgerblue",
            "firebrick","floralwhite","forestgreen","fuchsia","gainsboro","gold","goldenrod","gray","grey","green","greenyellow",
            "hotpink","indianred","indigo","khaki","lavender","lavenderblush","lawngreen","lemonchiffon","lightblue","lightcoral",
            "lightcyan","lightgoldenrodyellow","lightgray","lightgreen","lightgrey","lightpink","lightsalmon","lightseagreen",
            "lightskyblue","lightslategray","lightslategrey","lightsteelblue","lightyellow","limegreen","linen","magenta",
            "maroon","mediumaquamarine","mediumblue","mediumorchid","mediumpurple","mediumseagreen","mediumslateblue",
            "mediumspringgreen","mediumturquoise","mediumvioletred","midnightblue","mistyrose","moccasin","navy","oldlace","olive",
            "olivedrab","orange","orangered","orchid","palegoldenrod","palegreen","paleturquoise","palevioletred","papayawhip",
            "peachpuff","peru","pink","plum","powderblue","purple","red","rosybrown","royalblue","saddlebrown","salmon","sandybrown",
            "seagreen","sienna","silver","skyblue","slateblue","slategray","slategrey","springgreen","steelblue","tan","teal",
            "thistle","tomato","turquoise","violet","wheat","whitesmoke","yellow","yellowgreen" };*/

        auto foundIt = colorMapping.find(hash);
        if (foundIt == colorMapping.end())
        {
            auto& color = colorMapping[hash];
            color = nextColorIndex;
            nextColorIndex = (nextColorIndex + 1) % colorSchemeSize();
            return color;
        }
        else
        {
            return foundIt->second;
        }
    }
};

constexpr const char* singleIndent()
{
    return "    ";
}

inline std::string makeQuoted(const std::string& s)
{
    return "\"" + s + "\"";
}

inline std::string makeId(const std::string& path, INode* node)
{
    return path + "/" + node->getName().toString();
}

inline std::string makeClusterId(const std::string& id)
{
    return "Cluster//" + id;
}

inline void writeNodeProperties(INode* node, std::ostream& out)
{
    if (node->isRoot())
        out << "["
            << "label=\"\" "
            << "shape=point"
            << "]";
    else
        out << "["
            << "label=" << makeQuoted(node->getName().toString()) << "]";
}

inline void writeNode(const std::string& indent, const std::string& path, INode* node, std::ostream& out)
{
    out << indent << makeQuoted(makeId(path, node)) << " ";
    writeNodeProperties(node, out);
    out << ";" << std::endl;
}

inline void writeConnection(const std::string& indent, const std::string& path, INode* nodeA, INode* nodeB, std::ostream& out)
{
    auto graphA = nodeA->getNodeGraphDef();
    auto graphB = nodeB->getNodeGraphDef();

    std::string pathA = makeId(path, nodeA);
    std::string pathB = makeId(path, nodeB);

    std::string nodeAId = graphA ? makeId(pathA, graphA->getRoot()) : pathA;
    std::string nodeBId = graphB ? makeId(pathB, graphB->getRoot()) : pathB;

    out << indent;
    out << makeQuoted(nodeAId);
    out << " -> ";
    out << makeQuoted(nodeBId);

    if (graphA || graphB)
    {
        out << "[";
        if (graphA)
            out << "ltail=" << makeQuoted(makeClusterId(pathA));
        if (graphB)
            out << "lhead=" << makeQuoted(makeClusterId(pathB));
        out << "]";
    }

    out << ";" << std::endl;
}

inline void writeSubgraphProperties(
    const std::string& indent, INode* node, INodeGraphDef* graph, GraphState& state, std::ostream& out)
{
    const auto& nodeGraphDefName = graph->getName();
    std::string nodeGraphDefNameLabel =
        carb::fmt::format("{}({})", nodeGraphDefName.getString().c_str(), nodeGraphDefName.getHash());
    // this is nested NodeGraphDef
    if (node)
    {
        const auto& nodeName = node->getName();
        std::string nodeNameLabel = carb::fmt::format("{}({})", nodeName.getString().c_str(), node->getIndexInTopology());

        out << indent << "label = " << makeQuoted(nodeNameLabel + " | " + nodeGraphDefNameLabel) << std::endl;
    }
    // this is top level NodeGraphDef
    else
    {
        out << indent << "label = " << makeQuoted("EXECUTION GRAPH | " + nodeGraphDefNameLabel) << std::endl;
    }

    auto color = state.getColor(graph->getName().getHash());
    out << indent << "color = " << color << std::endl;
    out << indent << "node [color = " << color << "]" << std::endl;
    out << indent << "edge [color = " << color << "]" << std::endl;
}

inline void writeGraph(const std::string& indent,
                       const std::string& path,
                       INode* node,
                       INodeGraphDef* graph,
                       GraphState& state,
                       std::ostream& out)
{
    out << indent << "subgraph " << makeQuoted(makeClusterId(path)) << " {" << std::endl;

    std::string thisIndent = indent + singleIndent();

    writeSubgraphProperties(thisIndent, node, graph, state, out);

    // for readability, we first write nodes...
    writeNode(thisIndent, path, graph->getRoot(), out);
    traversal_dfs<VisitFirst>(graph->getRoot(),
                              [&out, &thisIndent, &path, &state](auto info, INode* prev, INode* curr)
                              {
                                  auto nodeGraph = curr->getNodeGraphDef();
                                  if (nodeGraph)
                                      writeGraph(thisIndent, makeId(path, curr), curr, nodeGraph, state, out);
                                  else
                                      writeNode(thisIndent, path, curr, out);

                                  info.continueVisit(curr);
                              });

    // ... and then we write connections
    traversal_dfs<VisitAll>(graph->getRoot(),
                            [&out, &thisIndent, &path](auto info, INode* prev, INode* curr)
                            {
                                writeConnection(thisIndent, path, prev, curr, out);

                                if (info.isFirstVisit()) // visit all edges, continue traversal on the first one
                                    info.continueVisit(curr);
                            });

    out << indent << "}" << std::endl;
}

inline void writeGraphProperties(const std::string& indent, std::ostream& out)
{
    out << indent << "compound=true" << std::endl;
    out << indent << "colorscheme=" << colorScheme() << std::endl;
    out << indent << "node [shape=circle style=filled fontcolor=white color=black colorscheme=" << colorScheme() << "]"
        << std::endl;
    out << indent << "edge [colorscheme=" << colorScheme() << "]" << std::endl;
    out << indent << "rankdir=LR" << std::endl;
    out << indent << "style=rounded" << std::endl;
}

#endif // DOXYGEN_BUILD

} // namespace detail

inline void writeFlattenedAsGraphviz(omni::core::ObjectParam<IGraph> inGraph, std::ostream& out)
{
    using namespace detail;

    out << "digraph ExecutionGraph {" << std::endl;
    {
        writeGraphProperties(singleIndent(), out);

        GraphState state;
        writeGraph(singleIndent(), "", nullptr, inGraph->getNodeGraphDef(), state, out);
    }
    out << "}" << std::endl;
}

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