# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
# pylint: disable=import-outside-toplevel
"""Tensor Expression Debug Display (TEDD), visualizing Tensor Expression"""
import html
import json
import warnings
from graphviz import Digraph
from graphviz import Source
import tvm

TVMDD_TABLE_BODY_WIDTH = 30
# Must match enum IterVarType defined in include/tvm/expr.h
ITERVAR_TYPE_STRING_MAP = {
    0: ('kDataPar', '#FFFFFF'),
    1: ('kThreadIndex', '#2980B9'),
    2: ('kCommReduce', '#FAD7A0'),
    3: ('kOrdered', '#D35400'),
    4: ('kOpaque', '#ABB2B9'),
    5: ('kUnrolled', '#D2B4DE'),
    6: ('kVectorized', '#AED6F1'),
    7: ('kParallelized', '#F5B7B1'),
    8: ('kTensorized', '#A9DFBF'),
}

PALETTE = {
    0: '#000000',
    1: '#922B21',
    2: '#76448A',
    3: '#1F618D',
    4: '#148F77',
    5: '#B7950B',
    6: '#AF601A',
    7: '#F5B7B1',
    8: '#A9DFBF',
}

PALETTE_SIZE = 9

def dom_path_to_string(dom_path, prefix=""):
    path_string = prefix
    for index in dom_path:
        path_string = path_string + '_' + str(index)
    return path_string


def insert_dot_id(sch):
    """Insert unique ID for each node in the DOM tree.
       They are used as Dot node ID.
       """
    for stage_idx, stage in enumerate(sch["stages"]):
        dom_path = [stage_idx]
        stage["id"] = dom_path_to_string(dom_path, stage["type"])
        for itervar_idx, itervar in enumerate(stage["all_itervars"]):
            dom_path = [stage_idx, itervar_idx]
            itervar["id"] = dom_path_to_string(dom_path, itervar["type"])
        for rel_idx, rel in enumerate(stage["relations"]):
            dom_path = [stage_idx, rel_idx]
            rel["id"] = dom_path_to_string(dom_path, rel["type"])
        for tensor_idx, tensor in enumerate(stage["output_tensors"]):
            dom_path = [stage_idx, tensor_idx]
            tensor["id"] = dom_path_to_string(dom_path, tensor["type"])
    return sch


class ObjectManager:
    """A helper class tracking schedule objects, e.g. stage, IterVar,
       relationship, and tensor, to their DOM path."""
    def __init__(self, sch):
        self.dict = {}
        for stage_idx, stage in enumerate(sch.stages):
            self.dict[stage] = [stage_idx]
            for itervar_idx, itervar in enumerate(stage.all_iter_vars):
                self.dict[itervar] = [stage_idx, itervar_idx]
            for rel_idx, rel in enumerate(stage.relations):
                self.dict[rel] = [stage_idx, rel_idx]
            for tensor_idx in range(stage.op.num_outputs):
                self.dict[frozenset({stage.op.name,
                                     tensor_idx})] = [stage_idx, tensor_idx]

    def get_dom_path(self, obj):
        if obj is None:
            return None
        assert obj in self.dict, 'Node is no found.'
        return self.dict[obj]


def get_or_create_dot_id(obj, prefix="", assert_on_missing=False):
    """If obj's ID has been registered, return it.
       If not, either assert or create a unique and legal ID, register and
       return it, according to assert_on_missing.
       ID must be a unique and legal Dotty ID.

        Parameters
        ----------
        obj : objet
                    Serve as the key to the ID.

        prefix : string
                    Prefix to attach to the ID.  Usually use obj's non-unique
                    name as prefix.

        assert_on_missing : bool
                    Assert or not if object doesn't have a registered ID.
    """
    prefix = prefix.replace('.', '_')
    if not hasattr(get_or_create_dot_id, "obj_id_dict"):
        get_or_create_dot_id.obj_id_dict = {}
    if obj not in get_or_create_dot_id.obj_id_dict:
        if assert_on_missing:
            assert False, 'dot_id ' + str(obj) + ' has not been registered.'
        else:
            get_or_create_dot_id.obj_id_dict[obj] = prefix + hex(id(obj))
    return get_or_create_dot_id.obj_id_dict[obj]


def get_port_id(is_input, index):
    return 'I_' + str(index) if is_input else 'O_' + str(index)


def get_itervar_type_info(iter_type):
    assert iter_type < len(
        ITERVAR_TYPE_STRING_MAP), 'Unknown IterVar type: ' + str(iter_type)
    return ITERVAR_TYPE_STRING_MAP[iter_type]


def get_itervar_label_color(itervar, iv_type):
    type_info = get_itervar_type_info(iv_type)
    return linebrk(
        str(itervar["name"]) + '(' + type_info[0] + ')',
        TVMDD_TABLE_BODY_WIDTH), type_info[1]


def linebrk(s, n):
    """ Break input string s with <br/> for every n charactors."""
    result = ''
    j = 0
    for i, c in enumerate(s):
        if j == n and i != len(s) - 1:
            result = result + '\n'
            j = 0
        j = j + 1
        result = result + c
    result = html.escape(str(result), quote=True)
    result = result.replace('\n', '<br/>')
    return result


def create_graph(name="", rankdir='BT'):
    graph = Digraph(name=name)
    graph.graph_attr['rankdir'] = rankdir
    return graph


def itervar_label(itervar, index, index_color, label):
    return '<TR><TD PORT="' + itervar[
        "id"] + '" BGCOLOR="' + index_color + '">' + str(
            index
        ) + '</TD><TD BGCOLOR="white" PORT="itervar">' + label + '<br/>' + str(
            itervar["properties"]["range"]) + '</TD></TR>'


def stage_label(stage):
    return stage['name'] + '<br/>Scope: ' + stage['properties']['scope']


def legend_label():
    label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" CELLPADDING="4">'
    for iter_type in ITERVAR_TYPE_STRING_MAP:
        name, color = ITERVAR_TYPE_STRING_MAP[iter_type]
        label += '<TR><TD BGCOLOR="' + color + '"></TD>' \
            + '<TD BGCOLOR="white">' + name + '</TD></TR>'
    label += '</TABLE>>'
    return label


def leaf_itervars(stage):
    filtered = filter(lambda x: (x["index"] >= 0), stage["all_itervars"])
    return sorted(filtered, key=lambda x: x["index"])


def legend_dot(g):
    with g.subgraph(name='cluster_legend') as subgraph:
        subgraph.attr(label='Legend')
        label = legend_label()
        subgraph.node('legend', label, shape='none', margin='0')


def extract_dom_for_viz(sch, need_range=True):
    json_str = dump_json(sch, need_range)
    s = json.loads(json_str)
    s = insert_dot_id(s)
    return s


def dump_graph(dot_string,
               show_svg=True,
               dot_file_path='',
               output_dot_string=False):
    """Output dot_string in various formats."""
    if dot_file_path:
        try:
            dot_file = open(dot_file_path, "w+")
            dot_file.write(dot_string)
            dot_file.close()
        except IOError:
            print('Cannot open file: ' + dot_file_path)
    if show_svg:
        from IPython.display import display
        from IPython.display import SVG
        src = Source(dot_string)
        display(SVG(src.pipe(format='svg')))
    if output_dot_string:
        return dot_string
    return None


def dump_json(sch, need_range):
    """Serialize data for visualization from a schedule in JSON format.

        Parameters
        ----------
        sch : schedule
                    The schedule object to serialize

        Returns
        -------
        json : string
            Serialized JSON string
    """
    def encode_itervar(itervar, stage, index, range_map):
        """Extract and encode IterVar visualization data to a dictionary"""
        ivrange = range_map[
            itervar] if range_map is not None and itervar in range_map else None
        bind_thread = None
        tensor_intrin = None
        if itervar in stage.iter_var_attrs:
            attr = stage.iter_var_attrs[itervar]
            iv_type = attr.iter_type
            # binding
            bind_thread = str(
                attr.bind_thread.var) if attr.bind_thread is not None else None
            # tensorization
            if attr.tensor_intrin is not None:
                tensor_intrin = str(attr.tensor_intrin.body)
                # remove the final \n
                tensor_intrin = tensor_intrin[0:-1] if tensor_intrin[
                    -1] == "\n" else tensor_intrin
            else:
                tensor_intrin = None
        else:
            iv_type = itervar.iter_type
        itervar_dict = {
            "type": "IterVar",
            "index": index,
            "name": str(itervar.var),
            "itervar_type": iv_type,
            "properties": {
                "thread": bind_thread,
                "intrin": tensor_intrin,
                "range": str(ivrange) if ivrange is not None else 'range(N/A)',
            }
        }
        return itervar_dict

    def encode_itervars(stage, range_map):
        """Extract and encode IterVars visualization data from a stage to a dictionary"""
        def get_leaf_itervar_index(itervar, leaf_iv):
            for leaf_index, ivar in enumerate(leaf_iv):
                if ivar == itervar:
                    return leaf_index
            return -1

        itervars = []
        for itervar in stage.all_iter_vars:
            leaf_index = get_leaf_itervar_index(itervar, stage.leaf_iter_vars)
            itervars.append(
                encode_itervar(itervar, stage, leaf_index, range_map))
        return itervars

    def encode_itervar_relation(obj_manager, rel):
        """Extract and encode IterVar Relationship visualization data to a dictionary"""
        rel_type = type(rel)
        if rel_type is tvm.te.schedule.Split:
            node_type = 'Split_Relation'
            rel_dict = {
                "type": node_type,
                "parent": obj_manager.get_dom_path(rel.parent),
                "outer": obj_manager.get_dom_path(rel.outer),
                "inner": obj_manager.get_dom_path(rel.inner),
            }
        elif rel_type is tvm.te.schedule.Fuse:
            node_type = 'Fuse_Relation'
            rel_dict = {
                "type": node_type,
                "fused": obj_manager.get_dom_path(rel.fused),
                "outer": obj_manager.get_dom_path(rel.outer),
                "inner": obj_manager.get_dom_path(rel.inner),
            }
        elif rel_type is tvm.te.schedule.Singleton:
            node_type = 'Singleton_Relation'
            rel_dict = {
                "type": node_type,
                "iter": obj_manager.get_dom_path(rel.iter),
            }
        else:
            return None
        return rel_dict

    def encode_itervar_relations(obj_manager, stage):
        relations = []
        for i in range(len(stage.relations)):
            rel = encode_itervar_relation(obj_manager, stage.relations[i])
            if rel is not None:
                relations.append(rel)
        return relations

    def encode_tensor(obj_manager, tensor, stage):
        """Extract and encode tensor visualization data to a dictionary"""
        tensor_dict = {
            "type": "Tensor",
            "source": obj_manager.get_dom_path(stage),
            "value_index": tensor.value_index,
            "shape": str(tensor.op.output(tensor.value_index).shape),
            "data_type": tensor.op.output(tensor.value_index).dtype,
        }
        return tensor_dict

    def encode_tensors(obj_manager, stage):
        tensors = []
        for i in range(stage.op.num_outputs):
            tensor = stage.op.output(i)
            tensors.append(encode_tensor(obj_manager, tensor, stage))
        tensors.sort(key=lambda tensor: tensor["value_index"])
        return tensors

    def encode_stage(obj_manager, stage, range_map):
        """Extract and encode stage visualization data to a dictionary"""
        stage_dict = {
            "type":
            "Stage",
            "name":
            stage.op.name,
            "attaching_to":
            obj_manager.get_dom_path(stage.attach_ivar),
            "compute":
            str(stage.op.body) if hasattr(stage.op, 'body') else None,
            "properties": {
                "scope": stage.scope,
            },
            "all_itervars":
            encode_itervars(stage, range_map),
            "relations":
            encode_itervar_relations(obj_manager, stage),
            "input_tensors": [
                obj_manager.get_dom_path(
                    frozenset({tensor.op.name, tensor.value_index}))
                for tensor in stage.op.input_tensors
            ],
            "output_tensors":
            encode_tensors(obj_manager, stage),
        }
        return stage_dict

    def encode_schedule(sch, need_range):
        """Extract and encode data from a schedule for visualization to a nested dictionary.
        It is useful for JSON to serialize schedule.

            Parameters
            ----------
            sch : schedule
                        The schedule object to extract

            Returns
            -------
            dict : dictionary
                A nested dictionary
        """
        assert isinstance(sch, tvm.te.schedule.Schedule
                          ), 'Input is not a tvm.te.schedule.Schedule object.'
        range_map = None
        if need_range:
            try:
                range_map = tvm.te.schedule.InferBound(sch)
            except tvm._ffi.base.TVMError as expt:
                warnings.warn(
                    'Ranges are not available, because InferBound fails with the following error:\n'
                    + str(expt))

        obj_manager = ObjectManager(sch)
        stages = []
        for stage in sch.stages:
            stages.append(encode_stage(obj_manager, stage, range_map))
        return {
            "type": "Schedule",
            "stages": stages,
        }

    return json.dumps(sch, default=lambda s: encode_schedule(s, need_range))


def viz_schedule_tree(sch,
                      show_svg=False,
                      dot_file_path='',
                      output_dot_string=False):
    """Top level API to render schedule tree

        Parameters
        ----------
        sch : schedule
                    The schedule object to visualize

        show_svg : bool
                    Display graph as SVG, useful for Jupyter notebooks.

        dot_file_path : string
                    Dot file to save the graph.

        output_dot_string : bool
                    Return dot file content or an empty string.

        Returns
        -------
        dot_string : string
            Dot file content or an empty string according to output_dot_string

        Examples
        --------
        The following code writes a schedule tree to a dot file.

        .. code-block:: python
            tedd.viz_schedule_tree(s, dot_file_path = '/tmp/example.dot')

        Use the following code to render a SVG graph in a Jupyter notebook.

        .. code-block:: python
            tedd.viz_schedule_tree(s, show_svg = True)
    """
    def create_schedule_tree_graph(name=""):
        return create_graph(name=name, rankdir='BT')

    def root_dot(g):
        g.node('ROOT', 'ROOT', shape='oval', margin='0')

    def stage_node_dot(g, stage):
        node_label = stage_node_label(stage)
        g.node(stage['id'], node_label, shape='none', margin='0')

    def stage_node_label(stage):
        """Return a html format label for the given stage."""
        label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" ' \
            'CELLPADDING="4"> <TR><TD BGCOLOR="lightgrey" ' \
            'COLSPAN="2" PORT="stage">' + stage_label(stage) + '</TD></TR>'

        for leafiv in leaf_itervars(stage):
            iv_type = leafiv["itervar_type"]
            var_attr_label = ''
            if "thread" in leafiv["properties"] and \
                    leafiv["properties"]["thread"] is not None:
                var_attr_label = var_attr_label + "<br/><font color=\"#2980B9\">(" + str(
                    leafiv["properties"]["thread"]) + ")</font>"
            if "intrin" in leafiv["properties"] and \
                    leafiv["properties"]["intrin"] is not None:
                var_attr_label = var_attr_label + "<br/>" + \
                    linebrk("(tensor_intrin:" + str(
                        leafiv["properties"]["intrin"]) + ")", TVMDD_TABLE_BODY_WIDTH)
            var_label, color = get_itervar_label_color(leafiv, iv_type)
            label += itervar_label(leafiv, leafiv["index"], color,
                                   var_label + var_attr_label)
        if stage["compute"] is not None:
            label += '<TR><TD COLSPAN="2">' + linebrk(str(
                stage["compute"]), TVMDD_TABLE_BODY_WIDTH) + '</TD></TR>'
        label += '</TABLE>>'
        return label

    def compute_at_dot(g, stage):
        """If the given stage attaches to another stage, create an edge from it
        stage to its attach point; otherwise, create an edge to the ROOT.
        """
        src = stage["id"]
        dst = dom_path_to_string(
            [stage["attaching_to"][0]], "Stage") + ":" + dom_path_to_string(
                stage["attaching_to"],
                "IterVar") if stage["attaching_to"] is not None else "ROOT"
        color = PALETTE[
            stage["attaching_to"][1] +
            1] if stage["attaching_to"] is not None and stage["attaching_to"][
                1] < PALETTE_SIZE - 1 else PALETTE[0]
        g.edge(src, dst, color=color)

    graph = create_schedule_tree_graph("Schedule Tree")
    s = extract_dom_for_viz(sch)
    legend_dot(graph)
    for stage in s['stages']:
        stage_node_dot(graph, stage)
    for stage in s['stages']:
        compute_at_dot(graph, stage)
    root_dot(graph)
    return dump_graph(graph.source, show_svg, dot_file_path, output_dot_string)


def viz_itervar_relationship_graph(sch,
                                   show_svg=False,
                                   dot_file_path='',
                                   output_dot_string=False):
    """Top level API to render IterVar relationship graph

        Parameters
        ----------
        sch : schedule
                    The schedule object to visualize

        show_svg : bool
                    Display graph as SVG, useful for Jupyter notebooks.

        dot_file_path : string
                    Dot file to save the graph.

        output_dot_string : bool
                    Return dot file content or an empty string.

        Examples
        --------
        The following code writes Ian tervar relationship graph to a dot file.

        .. code-block:: python
            tedd.viz_def viz_itervar_relationship_graph(sch,
                (s, dot_file_path = '/tmp/example.dot')

        Use the following code to render a SVG graph in a Jupyter notebook.

        .. code-block:: python
            tedd.viz_def viz_itervar_relationship_graph(sch,
                (s, show_svg = True)
    """
    def create_itervar_relation_graph(name=""):
        return create_graph(name=name, rankdir='TB')

    def itervar_node_dot(g, itervar, iv_type, index):
        label = itervar_node_label(itervar, iv_type, index)
        g.node(itervar["id"], label, shape='none', margin='0')

    def itervar_node_label(itervar, iv_type, index):
        label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" ' \
            'CELLPADDING="4">' + itervar_label(
                itervar, index,
                get_itervar_label_color(itervar, iv_type)[1],
                get_itervar_label_color(itervar, iv_type)[0]) + '</TABLE>>'
        return label

    def itervar_relation_node_dot(g, node_id, node_label, input_ports,
                                  output_ports):
        label = itervar_relation_node_label(node_label, input_ports,
                                            output_ports)
        g.node(node_id, label, shape='none', margin='0')

    def itervar_relation_node_label(node_label, input_ports, output_ports):
        """Return a html format label for an itervar relationship node
        including node_label and input/output ports.
        """
        label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" ' \
            'CELLPADDING="4">' + '<TR>'
        max_port_num = max(len(input_ports), len(output_ports))
        for i in range(max_port_num):
            if i < len(input_ports):
                input_port = input_ports[i]
                label += '<TD BGCOLOR="lightgrey" PORT="' + input_port + '">' \
                    + input_port + '</TD>'
            else:
                label += '<TD BGCOLOR="white"></TD>'
        label += '</TR>'
        label += '<TR><TD BGCOLOR="white" COLSPAN="' + str(
            max_port_num) + '" PORT="relation">' + node_label + '</TD></TR>'
        label += '<TR>'
        for i in range(max_port_num):
            if i < len(output_ports):
                output_port = output_ports[i]
                label += '<TD BGCOLOR="lightgrey" PORT="' + output_port + '">' \
                    + output_port + '</TD>'
            else:
                label += '<TD BGCOLOR="white"></TD>'
        label += '</TR>'
        label += '</TABLE>>'
        return label

    def itervar_relation_dot(g, node, node_id):
        """Create an itervar relationship node."""
        node_type = node["type"]
        if node_type == "Split_Relation":
            node_type = 'Split'
            itervar_relation_node_dot(g, node_id, node_type, ['Input'],
                                      ['Outer', 'Inner'])
            parent = dom_path_to_string(node["parent"], "IterVar")
            outer = dom_path_to_string(node["outer"], "IterVar")
            inner = dom_path_to_string(node["inner"], "IterVar")
            g.edge(parent + ':itervar', node_id + ':Input')
            g.edge(node_id + ':Outer', outer + ':itervar')
            g.edge(node_id + ':Inner', inner + ':itervar')
        elif node_type == "Fuse_Relation":
            node_type = 'Fuse'
            itervar_relation_node_dot(g, node_id, node_type,
                                      ['Outer', 'Inner'], ['Fused'])
            fused = dom_path_to_string(node["fused"], "IterVar")
            outer = dom_path_to_string(node["outer"], "IterVar")
            inner = dom_path_to_string(node["inner"], "IterVar")
            g.edge(outer + ':itervar', node_id + ':Outer')
            g.edge(inner + ':itervar', node_id + ':Inner')
            g.edge(node_id + ':Fused', fused + ':itervar')
        elif node_type == "Singleton_Relation":
            node_type = 'Singleton'
            itervar_relation_node_dot(g, node_id, node_type, [], ['Iter'])
            itervar = dom_path_to_string(node["inner"], "IterVar")
            g.edge(node_id + ':Iter', itervar + ':itervar')
        else:
            assert False, 'Unknown IterVarRelationNode: ' + node_type

    def stage_node_dot(g, stage):
        """Create a stage node."""
        with g.subgraph(name='cluster_' + stage["id"]) as subgraph:
            subgraph.attr(label=stage["name"])
            if stage["all_itervars"]:
                for itervar in stage["all_itervars"]:
                    iv_type = itervar["itervar_type"]
                    itervar_node_dot(subgraph, itervar, iv_type,
                                     itervar["index"])
                for rel in stage["relations"]:
                    node_id = rel["id"]
                    itervar_relation_dot(subgraph, rel, node_id)
            else:
                subgraph.node(stage["name"] + '_placeholder', style='invis')

    graph = create_itervar_relation_graph("IterVar Relationship Graph")
    s = extract_dom_for_viz(sch)
    legend_dot(graph)
    for stage in s['stages']:
        stage_node_dot(graph, stage)

    return dump_graph(graph.source, show_svg, dot_file_path, output_dot_string)


def viz_dataflow_graph(sch,
                       show_svg=False,
                       dot_file_path='',
                       output_dot_string=False):
    """Top level API to render dataflow graph

        Parameters
        ----------
        sch : schedule
                    The schedule object to visualize

        show_svg : bool
                    Display graph as SVG, useful for Jupyter notebooks.

        dot_file_path : string
                    Dot file to save the graph.

        output_dot_string : bool
                    Return dot file content or an empty string.

        Examples
        --------
        The following code writes a dataflow graph to a dot file.

        .. code-block:: python
            tedd.viz_dataflow_graph(s, dot_file_path = '/tmp/example.dot')

        Use the following code to render a SVG graph in a Jupyter notebook.

        .. code-block:: python
            tedd.viz_dataflow_graph(s, show_svg = True)    """
    def create_dataflow_graph(name=""):
        return create_graph(name=name, rankdir='LR')

    def tensor_node_dot(g, tensor):
        """Create a tensor node."""
        label = tensor_node_label(tensor)
        g.node(tensor["id"], label, shape='oval', margin='0')

    def tensor_node_label(tensor):
        """Return a html format label for the given tensor."""
        label = str(tensor["shape"]) + '\n' + str(tensor["data_type"])
        return label

    def stage_node_dot(g, stage):
        """Create a stage node."""
        label = stage_node_label(stage)
        g.node(stage["id"], label, shape='none', margin='0')

    def stage_node_label(stage):
        """Return a html format label for the given stage."""
        rows = max(
            1, max(len(stage["output_tensors"]), len(stage["input_tensors"])))
        label = '<<TABLE BORDER="0" CELLBORDER="1" CELLSPACING="0" ' \
            'CELLPADDING="4">'
        for i in range(rows):
            label += '<TR>'
            if i < len(stage["input_tensors"]):
                port_id = get_port_id(True, i)
                label += '<TD BGCOLOR="lightgrey" COLSPAN="2" PORT="' \
                    + port_id + '">' + str(
                        i) + '</TD>'
            else:
                label += '<TD BGCOLOR="white" COLSPAN="2"></TD>'
            if i == 0:
                label += '<TD BGCOLOR="white" COLSPAN="2" ROWSPAN="' + str(
                    rows) + '">' + stage_label(stage) + '</TD>'
            if i < len(stage["output_tensors"]):
                port_id = get_port_id(False, i)
                label += '<TD BGCOLOR="lightgrey" COLSPAN="2" PORT="' \
                    + port_id + '">' + str(
                        i) + '</TD>'
            else:
                label += '<TD BGCOLOR="white" COLSPAN="2"></TD>'
            label += '</TR>'
        label += '</TABLE>>'
        return label

    def dfg_dot(g, sch):
        """Create edges among stages."""
        stages = sch['stages']
        for stage in stages:
            for i in range(len(stage["input_tensors"])):
                src = dom_path_to_string(stage["input_tensors"][i], "Tensor")
                dst = stage["id"] + ':' + get_port_id(True, i)
                g.edge(src, dst)
            for i in range(len(stage["output_tensors"])):
                src = stage["id"] + ':' + get_port_id(False, i)
                dst = stage["output_tensors"][i]["id"]
                g.edge(src, dst)

    graph = create_dataflow_graph("Dataflow Graph")
    s = extract_dom_for_viz(sch, need_range=False)
    for stage in s['stages']:
        stage_node_dot(graph, stage)
        for tensor in stage["output_tensors"]:
            tensor_node_dot(graph, tensor)

    dfg_dot(graph, s)

    return dump_graph(graph.source, show_svg, dot_file_path, output_dot_string)
