"""
Experimental. Tools for visualizing the torch.jit.Graph objects.
"""
import string
import json

_vis_template = string.Template("""
<!doctype html>
<html>
<head>
  <title>$name</title>

  <script type="text/javascript" src="https://cdnjs.cloudflare.com/ajax/libs/vis/4.20.1/vis.min.js"></script>
  <link href="https://cdnjs.cloudflare.com/ajax/libs/vis/4.20.1/vis.min.css" rel="stylesheet" type="text/css" />

  <style type="text/css">
    #mynetwork {
      height: 100vh;
    }
  </style>
</head>
<body>

<div id="mynetwork"></div>

<script type="text/javascript">
  // create an array with nodes
  var nodes = new vis.DataSet(
    $nodes
  );

  // create an array with edges
  var edges = new vis.DataSet(
    $edges
  );

  // create a network
  var container = document.getElementById('mynetwork');
  var data = {
    nodes: nodes,
    edges: edges
  };
  var options = $options;
  var network = new vis.Network(container, data, options);
</script>
</body>
</html>
""")


def write(self, filename):
    """
    Write an html file that visualizes a torch.jit.Graph using vis.js
    Arguments:
        self (torch.jit.Graph): the graph.
        filename (string): the output filename, an html-file.
    """

    nodes = []
    edges = []
    options = {}
    for n, i in enumerate(self.inputs()):
        nodes.append({
            'id': i.unique(),
            'label': 'input {}'.format(n),
            'shape': 'square',
        })

    existing = set()

    def add_edge(i_, n):
        i = i_ if i_.kind() != 'Select' else i_.input()
        if (i, n) in existing:
            return
        existing.add((i, n))
        e = {
            'from': n.unique(),
            'to': i.unique(),
            'arrows': 'from',
        }
        if i.stage() != n.stage():
            e['color'] = 'green'
        edges.append(e)

    counts = {}
    offset = 0
    for n in self.nodes():
        if len(n.uses()) == 0 or n.kind() == 'Undefined':
            continue
        ident = counts.get(n.kind(), 0)
        counts[n.kind()] = ident + 1
        d = {
            'id': n.unique(),
            'label': '{}_{}'.format(n.kind(), ident),
            'y': offset,
            'fixed': {'y': True},
        }
        if n in self.outputs():
            d['shape'] = 'triangle'

        for i in n.inputs():
            add_edge(i, n)

        nodes.append(d)
        offset += 30

    result = _vis_template.substitute(nodes=json.dumps(nodes),
                                      edges=json.dumps(edges),
                                      options=json.dumps(options),
                                      name=filename)
    with open(filename, 'w') as f:
        f.write(result)
