/*******************************************************************************
* Copyright 2017-2018 Intel Corporation
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
*     http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*******************************************************************************/

#include <fstream>

#include "ngraph/function.hpp"
#include "ngraph/graph_util.hpp"
#include "ngraph/node.hpp"
#include "ngraph/pass/pass.hpp"
#include "ngraph/pass/visualize_tree.hpp"
#include "ngraph/util.hpp"

using namespace ngraph;
using namespace std;

bool pass::VisualizeTree::run_on_module(vector<shared_ptr<ngraph::Function>>& functions)
{
    for (shared_ptr<Function> f : functions)
    {
        // map<size_t, list<node_ptr>> dependent_nodes;
        traverse_nodes(f, [&](shared_ptr<Node> node) {
            for (auto arg : node->get_input_ops())
            {
                m_ss << add_attributes(arg);
                m_ss << add_attributes(node);
                m_ss << "    " << arg->get_name() << " -> " << node->get_name();
                m_ss << ";\n";
            }
        });
    }

    render();

    return false;
}

pass::VisualizeTree::VisualizeTree(const string& file_name)
    : m_name{file_name}
{
}

std::string pass::VisualizeTree::add_attributes(shared_ptr<Node> node)
{
    string rc;
    if (!contains(m_nodes_with_attributes, node))
    {
        m_nodes_with_attributes.insert(node);
        rc = get_attributes(node);
    }
    return rc;
}

std::string pass::VisualizeTree::get_attributes(shared_ptr<Node> node)
{
    stringstream ss;
    if (node->is_parameter())
    {
        ss << "    " << node->get_name() << " [shape=box color=blue";
    }
    else
    {
        ss << "    " << node->get_name() << " [shape=ellipse color=black";
    }

    ss << " label=\"" << node->get_name();

    if (std::getenv("NGRAPH_VISUALIZE_TREE_OUTPUT_SHAPES") != nullptr)
    {
        ss << " " << vector_to_string(node->get_shape());
    }

    ss << " \"]\n";

    return ss.str();
}

void pass::VisualizeTree::render() const
{
#ifdef GRAPHVIZ_FOUND
    auto tmp_file = m_name + ".tmp";
    ofstream out(tmp_file);
    if (out)
    {
        out << "digraph ngraph\n{\n";
        out << m_ss.str();
        out << "}\n";
        out.close();

        stringstream ss;

        const char* format = std::getenv("NGRAPH_VISUALIZE_TREE_OUTPUT_FORMAT");
        if (!format)
        {
            format = "png";
        }

        ss << "dot -T" << format << " " << tmp_file << " -o " << m_name;
        auto cmd = ss.str();
        auto stream = popen(cmd.c_str(), "r");
        pclose(stream);

        remove(tmp_file.c_str());
    }
#endif
}
