text: /**
   * @fileoverview Transformer Visualization D3 javascript code.
   *
   *
  
   *  Based on: https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/visualization/attention.js
   *
   * Change log:
   *
   * 12/19/18  Jesse Vig   Assorted cleanup. Changed orientation of attention matrices.
   */
  
  requirejs(['jquery', 'd3'], function($, d3) {
  
  const TEXT_SIZE = 15;
  const BOXWIDTH = TEXT_SIZE * 8;
  const BOXHEIGHT = TEXT_SIZE * 1.5;
  const WIDTH = 2000;
  const HEIGHT_PADDING = 100;
  const MATRIX_WIDTH = 150;
  const HEAD_COLORS = d3.scale.category10();
  const CHECKBOX_SIZE = 20;
  
  var params = window.params;
  var config = {};
  initialize();
  
  function lighten(color) {
    var c = d3.hsl(color);
    var increment = (1 - c.l) * 0.6;
    c.l += increment;
    c.s -= increment;
    return c;
  }
  
  function transpose(mat) {
    return mat[0].map(function(col, i) {
      return mat.map(function(row) {
        return row[i];
      });
    });
  }
  
  function zip(a, b) {
    return a.map(function (e, i) {
      return [e, b[i]];
    });
  }
  
  function render() {
  
    var attnData = config.attention[config.filter];
    var leftText = attnData.left_text;
    var rightText = attnData.right_text;
    var attentionHeads = attnData.attn[config.layer];
  
    $("#vis svg").empty();
    $("#vis").empty();
  
    var height = config.initialTextLength * BOXHEIGHT * 2 + HEIGHT_PADDING;
    var svg = d3.select("#vis")
              .append('svg')
              .attr("width", WIDTH)
              .attr("height", height);
  
    var attData = [];
    for (var i=0; i < config.nHeads; i++) {
      var att = attentionHeads[i];
      var att_trans = transpose(att);
      attData.push(zip(att_trans, att));
    }
  
    renderText(svg, leftText, true, attData, 0);
    renderText(svg, rightText, false, attData, MATRIX_WIDTH + BOXWIDTH);
  
    renderAttentionHighlights(svg, attData);
  
    svg.append("g").classed("attentionHeads", true);
  
    renderAttention(svg, attentionHeads);
  
    drawCheckboxes(0, svg, attentionHeads);
  
  }
  
  function renderText(svg, text, isLeft, attData, leftPos) {
    // attData: list of tuples (att, att_trans), one for each layer. att and att_trans are attention matrics for each layer.
    //           att is of shape [nHeads, source_len, target_len)
    var id = isLeft ? "left" : "right";
    var textContainer = svg.append("svg:g")
                           .attr("id", id);
  
    textContainer.append("g").classed("attentionBoxes", true)
                 .selectAll("g")
                 .data(attData)
                 .enter()
                 .append("g")
                 .selectAll("rect")
                 .data(function(d) {return d;})
                 .enter()
                 .append("rect")
                 .attr("x", function(d, i, j) {
                   return leftPos + boxOffsets(j);
                 })
                 .attr("y", function(d, i) {
                   return (+1) * BOXHEIGHT;
                 })
                 .attr("width", BOXWIDTH / activeHeads())
                 .attr("height", function() { return BOXHEIGHT; })
                 .attr("fill", function(d, i, j) {
                    return HEAD_COLORS(j);
                  })
                 .style("opacity", 0.0);
  
    var tokenContainer = textContainer.append("g").selectAll("g")
                                      .data(text)
                                      .enter()
                                      .append("g");
  
    tokenContainer.append("rect")
                  .classed("background", true)
                  .style("opacity", 0.0)
                  .attr("fill", "lightgray")
                  .attr("x", leftPos)
                  .attr("y", function(d, i) {
                    return (i+1) * BOXHEIGHT;
                  })
                  .attr("width", BOXWIDTH)
                  .attr("height", BOXHEIGHT);
  
    var textEl = tokenContainer.append("text")
                                .text(function(d) { return d; })
                                .attr("font-size", TEXT_SIZE + "px")
                                .style("cursor", "default")
                                .style("-webkit-user-select", "none")
                                .attr("x", leftPos)
                                .attr("y", function(d, i) {
                                  return (i+1) * BOXHEIGHT;
                                });
  
    if (isLeft) {
      textEl.style("text-anchor", "end")
             .attr("dx", BOXWIDTH - TEXT_SIZE)
             .attr("dy", TEXT_SIZE);
    } else {
      textEl.style("text-anchor", "start")
             .attr("dx", + TEXT_SIZE)
             .attr("dy", TEXT_SIZE);
    }
  
    tokenContainer.on("mouseover", function(d, index) {
      textContainer.selectAll(".background")
                   .style("opacity", function(d, i) {
                     return i == index ? 1.0 : 0.0;
                   });
  
      svg.selectAll(".attentionHeads").style("display", "none");
  
      svg.selectAll(".lineHeads")  // To get the nesting to work.
         .selectAll(".attLines")
         .attr("stroke-opacity", function(d) {
            return 1.0;
          })
         .attr("y1", function(d, i) {
          if (isLeft) {
            return (index+1) * BOXHEIGHT + (BOXHEIGHT/2);
          } else {
            return (i+1) * BOXHEIGHT + (BOXHEIGHT/2);
          }
       })
       .attr("x1", BOXWIDTH)
       .attr("y2", function(d, i) {
         if (isLeft) {
            return (i+1) * BOXHEIGHT + (BOXHEIGHT/2);
          } else {
            return (index+1) * BOXHEIGHT + (BOXHEIGHT/2);
          }
       })
       .attr("x2", BOXWIDTH + MATRIX_WIDTH)
       .attr("stroke-width", 2)
       .attr("stroke", function(d, i, j) {
          return HEAD_COLORS(j);
        })
       .attr("stroke-opacity", function(d, i, j) {
        if (isLeft) {d = d[0];} else {d = d[1];}
        if (config.headVis[j]) {
          if (d) {
            return d[index];
          } else {
            return 0.0;
          }
        } else {
          return 0.0;
        }
       });
  
      function updateAttentionBoxes() {
        var id = isLeft ? "right" : "left";
        var leftPos = isLeft ? MATRIX_WIDTH + BOXWIDTH : 0;
        svg.select("#" + id)
           .selectAll(".attentionBoxes")
           .selectAll("g")
           .selectAll("rect")
           .attr("x", function(d, i, j) { return leftPos + boxOffsets(j); })
           .attr("y", function(d, i) { return (i+1) * BOXHEIGHT; })
           .attr("width", BOXWIDTH/activeHeads())
           .attr("height", function() { return BOXHEIGHT; })
           .style("opacity", function(d, i, j) {
              if (isLeft) {d = d[0];} else {d = d[1];}
              if (config.headVis[j])
                if (d) {
                  return d[index];
                } else {
                  return 0.0;
                }
              else
                return 0.0;
           });
      }
  
      updateAttentionBoxes();
    });
  
    textContainer.on("mouseleave", function() {
      d3.select(this).selectAll(".background")
                     .style("opacity", 0.0);
      svg.selectAll(".attLines").attr("stroke-opacity", 0.0);
      svg.selectAll(".attentionHeads").style("display", "inline");
      svg.selectAll(".attentionBoxes")
         .selectAll("g")
         .selectAll("rect")
         .style("opacity", 0.0);
    });
  }
  
  function renderAttentionHighlights(svg, attention) {
    var line_container = svg.append("g");
    line_container.selectAll("g")
                  .data(attention)
                  .enter()
                  .append("g")
                  .classed("lineHeads", true)
                  .selectAll("line")
                  .data(function(d){return d;})
                  .enter()
                  .append("line").classed("attLines", true);
  }
  
  function renderAttention(svg, attentionHeads) {
    var line_container = svg.selectAll(".attentionHeads");
    line_container.html(null);
    for(var h=0; h<attentionHeads.length; h++) {
      for(var s=0; s<attentionHeads[h].length; s++) {
        for(var a=0; a<attentionHeads[h][s].length; a++) {
          line_container.append("line")
          .attr("y1", (s+1) * BOXHEIGHT + (BOXHEIGHT/2))
          .attr("x1", BOXWIDTH)
          .attr("y2", (a+1) * BOXHEIGHT + (BOXHEIGHT/2))
          .attr("x2", BOXWIDTH + MATRIX_WIDTH)
          .attr("stroke-width", 2)
          .attr("stroke", HEAD_COLORS(h))
          .attr("stroke-opacity", function() {
            if (config.headVis[h]) {
              return attentionHeads[h][s][a]/activeHeads();
            } else {
              return 0.0;
            }
          }());
        }
      }
    }
  }
  
  // Checkboxes
  function boxOffsets(i) {
    var numHeadsAbove = config.headVis.reduce(
        function(acc, val, cur) {return val && cur < i ? acc + 1: acc;}, 0);
    return numHeadsAbove * (BOXWIDTH / activeHeads());
  }
  
  function activeHeads() {
    return config.headVis.reduce(function(acc, val) {
      return val ? acc + 1: acc;
    }, 0);
  }
  
  function drawCheckboxes(top, svg, attentionHeads) {
    var checkboxContainer = svg.append("g");
    var checkbox = checkboxContainer.selectAll("rect")
                                    .data(config.headVis)
                                    .enter()
                                    .append("rect")
                                    .attr("fill", function(d, i) {
                                      return HEAD_COLORS(i);
                                    })
                                    .attr("x", function(d, i) {
                                      return (i+1) * CHECKBOX_SIZE;
                                    })
                                    .attr("y", top)
                                    .attr("width", CHECKBOX_SIZE)
                                    .attr("height", CHECKBOX_SIZE);
  
    function updateCheckboxes() {
      checkboxContainer.selectAll("rect")
                                .data(config.headVis)
                                .attr("fill", function(d, i) {
        var headColor = HEAD_COLORS(i);
        var color = d ? headColor : lighten(headColor);
        return color;
      });
    }
  
    updateCheckboxes();
  
    checkbox.on("click", function(d, i) {
      if (config.headVis[i] && activeHeads() == 1) return;
      config.headVis[i] = !config.headVis[i];
      updateCheckboxes();
      renderAttention(svg, attentionHeads);
    });
  
    checkbox.on("dblclick", function(d, i) {
      // If we double click on the only active head then reset
      if (config.headVis[i] && activeHeads() == 1) {
        config.headVis = new Array(config.nHeads).fill(true);
      } else {
        config.headVis = new Array(config.nHeads).fill(false);
        config.headVis[i] = true;
      }
      updateCheckboxes();
      renderAttention(svg, attentionHeads);
    });
  }
  
  function initialize() {
    config.attention = params['attention'];
    config.filter = params['default_filter'];
    config.nLayers = config.attention[config.filter]['attn'].length;
    console.log('num layers')
    console.log(config.nLayers)
    config.nHeads = config.attention[config.filter]['attn'][0].length;
    config.headVis  = new Array(config.nHeads).fill(true);
    config.layer = 0;
    config.initialTextLength = config.attention[config.filter].right_text.length
  }
  
  $("#layer").empty();
  for(var i=0; i<config.nLayers; i++) {
    $("#layer").append($("<option />").val(i).text(i));
  }
  
  $("#layer").on('change', function(e) {
    config.layer = +e.currentTarget.value;
    render();
  });
  
  $("#filter").on('change', function(e) {
    config.filter = e.currentTarget.value;
    render();
  });
  
  render();
  
  });
  