<!doctype html>
<!--
 Copyright 2019 Google LLC

 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

      https://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.
-->

<html lang="en-us">
  <head>
    <meta charset="utf-8">
    <title>IREE - MNIST Demo</title>
    <style>
      body {
        margin: 0;
        background-color: #444444;
      }

      p {
        margin: 4px;
      }

      .drop-target {
        border: 3px solid blue;
        background-color: #AAAAAA;
        width:  300px;
        height: 200px;
        margin: 4px;
        padding: 4px;
      }
    </style>
  </head>

  <body>
    <div id="drop-zone" class="drop-target" style="display:none">
      <p>Drag an IREE module (.fb) file here to load it</p>
    </div>

    <!-- Overview -->
    <div style="background-color: #AAAAAA; padding: 16px; margin: 16px; width: 800px">
      <div>
        <h2>IREE MNIST demo, CPU interpreter using WebAssembly</h2>
        <p>WebAssembly threads must be enabled. See chrome://flags/#enable-webassembly-threads</p>
      </div>

      <br>
      <hr>
      <br>

      <div style="display:none">
        <p>Input arguments:</p>
        <textarea id="input-arguments" name="arguments" placeholder="4x4xi8=0,1,2,3" style="width: 600px; box-sizing: border-box; padding: 16px;">1x5xf32=1,-2,-3,4,-5\n1x5x3x1xf32=15,14,13,12,11,10,9,8,7,6,5,4,3,2,1</textarea>
        <hr>
      </div>

      <p>Input image for MNIST (handwritten digit) classification:</p>
      <ul>
        <li>Left click drag to draw</li>
        <li>Right click drag to erase</li>
      </ul>
      <br>
      <canvas id="draw-canvas" style="width: 448px; height: 448px; border:1px solid #000000;"></canvas>
      <canvas id="zoom-canvas" style="width: 28px; height:28px; border:1px solid #000000;"></canvas>
      <br>
      <br>
      <button id="reset-canvas-button">Reset canvas</button>

      <br>
      <br>
      <hr><br>

      <div style="display:none">
        <button id="run-module-button">Run module</button>
        <br>
        <br>
      </div>
      <p style="display:inline-block">MNIST Prediction:&nbsp;
        <div id="prediction-summary" style="display:inline-block"></div>
      </p>
      <p>MNIST Prediction confidences:&nbsp;
        <p id="prediction-details">
          <br><br><br><br><br><br><br><br><br><br>
        </p>
      </p>
    </div>

    <script type="text/javascript">
      let isModuleInitialized = false;
      let canvasHasChanged = false;

      var Module = {
        onRuntimeInitialized: () => {
          console.log('onRuntimeInitialized');
          isModuleInitialized = true;
        },
        print: () => {
          return function(text) {
            console.log(Array.prototype.slice.call(arguments).join(' '));
          };
        },
        printErr: function(text) {
          console.error(Array.prototype.slice.call(arguments).join(' '));
        },
      };
      window.onerror = function() {
        console.log('onerror: ' + event);
      };

      function runModule(moduleData, moduleInputs) {
        return Module.runIreeModule(moduleData, moduleInputs);
      }

      // TODO(scotttodd): generalize this
      function getMnistPredictions(resultsString) {
        const resultsArrayString = resultsString.slice(10, resultsString.length - 1);
        const resultsArray = resultsArrayString.split(' ').map(parseFloat);
        return resultsArray;
      }

      function runLoadedModuleOnCanvasData() {
        let canvasArguments = '1x28x28x1xf32=';
        const zoomCtxData = zoomCtx.getImageData(
            0, 0, zoomCanvas.width, zoomCanvas.height).data;
        // Extract grayscale floats [0.0, 1.0] from uint8_t rgba triples
        for (let i = 0; i < zoomCtxData.length / 4; ++i) {
          const r = zoomCtxData[i * 4 + 0] / 255.0;
          const g = zoomCtxData[i * 4 + 1] / 255.0;
          const b = zoomCtxData[i * 4 + 2] / 255.0;
          const luma = 1.0 - (r * 0.299 + g * 0.587 + b * 0.114);
          if (i == zoomCtxData.length / 4 - 1) {
            canvasArguments += luma.toFixed(3);
          } else {
            canvasArguments += luma.toFixed(3) + ', ';
          }
        }

        // TODO(scotttodd): Validate loaded
        const resultsString = runModule(moduleFileData, canvasArguments);
        predictions = getMnistPredictions(resultsString);

        // Index with the highest confidence is the "final" prediction.
        let highestConfidencePrediction = 0.0;
        let highestConfidencePredictionIndex = 0;
        for (let i = 0; i < predictions.length; ++i) {
          if (predictions[i] > highestConfidencePrediction) {
            highestConfidencePrediction = predictions[i];
            highestConfidencePredictionIndex = i;
          }
        }
        predictionSummaryDiv.innerHTML = highestConfidencePredictionIndex;

        let detailsString = '';
        for (let i = 0; i < predictions.length; ++i) {
          detailsString += i + ': ' + (predictions[i] * 100).toFixed(2) + '%<br>';
        }
        predictionDetailsDiv.innerHTML = detailsString;
      }

      // <textarea> argument input
      const inputArguments = document.getElementById('input-arguments');

      const predictionSummaryDiv = document.getElementById('prediction-summary');
      const predictionDetailsDiv = document.getElementById('prediction-details');

      let moduleFileData = '';

      // <canvas> input for drawing (2d) images
      // Full size canvas for user input
      const drawCanvas = document.getElementById('draw-canvas');
      drawCanvas.width = 448;
      drawCanvas.height = 448;
      const drawCtx = drawCanvas.getContext('2d');
      // Zoomed canvas for sending data down into wasm
      const zoomCanvas = document.getElementById('zoom-canvas');
      zoomCanvas.width = 28;
      zoomCanvas.height = 28;
      const zoomCtx = zoomCanvas.getContext('2d');
      function clearCanvas(canvas, ctx) {
        ctx.fillStyle = 'white';
        ctx.fillRect(0, 0, canvas.width, canvas.height);
      }
      clearCanvas(drawCanvas, drawCtx);
      clearCanvas(zoomCanvas, zoomCtx);
      // User input handling (drawing, copying to the zoomed canvas)
      const mousePosition = { x: 0, y : 0};
      drawCanvas.addEventListener('mousedown', (event) => {
        mousePosition.x = event.offsetX;
        mousePosition.y = event.offsetY;
      });
      drawCanvas.addEventListener('mousemove', (event) => {
        if (event.buttons !== 1 && event.buttons !== 2) {
          return;
        }
        drawCtx.beginPath();
        drawCtx.lineWidth = 60;
        drawCtx.lineCap = 'round';
        drawCtx.strokeStyle = event.buttons === 1 ? '#000000' : '#FFFFFF';
        drawCtx.moveTo(mousePosition.x, mousePosition.y);
        mousePosition.x = event.offsetX;
        mousePosition.y = event.offsetY;
        drawCtx.lineTo(mousePosition.x, mousePosition.y);
        drawCtx.stroke();

        zoomCtx.drawImage(drawCanvas, 0, 0, zoomCanvas.width, zoomCanvas.height);

        canvasHasChanged = true;
      });
      drawCanvas.addEventListener('contextmenu', (event) => {
        event.preventDefault();
      });
      const resetCanvasButton = document.getElementById('reset-canvas-button');
      resetCanvasButton.addEventListener('click', () => {
        clearCanvas(drawCanvas, drawCtx);
        clearCanvas(zoomCanvas, zoomCtx);

        canvasHasChanged = true;
      });

      // IREE module loading from drag and dropped files
      const dropZone = document.getElementById('drop-zone');
      dropZone.addEventListener('drop', (dropEvent) => {
        dropEvent.preventDefault();
        dropEvent.target.style.border = '';

        console.log('File received, loading...');

        // Assume exactly one file was dropped.
        const uploadedFile = dropEvent.dataTransfer.items[0].getAsFile();
        const fileReader = new FileReader();
        fileReader.onload = (fileLoadEvent) => {
          console.log('File loaded, running IREE module...');
          // TODO(scotttodd): Wait until after onRuntimeInitialized() is called
          // Module.runIreeModule(fileLoadEvent.target.result, inputArguments.value);

          let canvasArguments = '1x28x28x1xf32=';
          const zoomCtxData = zoomCtx.getImageData(
              0, 0, zoomCanvas.width, zoomCanvas.height).data;
          // Extract grayscale floats [0.0, 1.0] from uint8_t rgba triples
          for (let i = 0; i < zoomCtxData.length / 4; ++i) {
            const r = zoomCtxData[i * 4 + 0] / 255.0;
            const g = zoomCtxData[i * 4 + 1] / 255.0;
            const b = zoomCtxData[i * 4 + 2] / 255.0;
            const luma = 1.0 - (r * 0.299 + g * 0.587 + b * 0.114);
            if (i == zoomCtxData.length / 4 - 1) {
              canvasArguments += luma.toFixed(3);
            } else {
              canvasArguments += luma.toFixed(3) + ', ';
            }
          }
          const result = runModule(fileLoadEvent.target.result, canvasArguments);
          console.log('result', result);
        };
        fileReader.readAsArrayBuffer(uploadedFile);
      });
      dropZone.addEventListener('dragover', (event) => {
        event.preventDefault();
      });
      dropZone.addEventListener('dragenter', (event) => {
        event.target.style.border = '3px dotted red';
      });
      dropZone.addEventListener('dragleave', (event) => {
        event.target.style.border = '';
      });

      // Fetch IREE module from the ?module=[foo.fb] URL query parameter,
      // if provided. Otherwise rely on drag-and-drop to get executable module.
      const urlParams = new URLSearchParams(window.location.search);
      if (urlParams.has('module')) {
        const modulePath = urlParams.get('module');
        console.log('modulePath', modulePath);

        const moduleRequest = new XMLHttpRequest();
        moduleRequest.responseType = 'arraybuffer';
        moduleRequest.onreadystatechange = () => {
          if (moduleRequest.readyState == XMLHttpRequest.DONE) {
            if (moduleRequest.status == 200) {
              console.log('Fetched module from ', modulePath);
              moduleFileData = moduleRequest.response;
            } else {
              console.error('Module request failed: ', moduleRequest.status);
            }
          }
        };
        moduleRequest.open("GET", modulePath, true);
        moduleRequest.send();
      } else {
        console.log('No module set in URL params, use drag-and-drop');
      }

      const runModuleButton = document.getElementById('run-module-button');
      runModuleButton.addEventListener('click', () => {
        runLoadedModuleOnCanvasData();
      });

      window.setInterval(() => {
        if (!isModuleInitialized || moduleFileData.length === 0 ||
            !canvasHasChanged) {
          return;
        }

        runLoadedModuleOnCanvasData();
        canvasHasChanged = false;
      }, 200);
    </script>
    <script src="run_module_emscripten.js"></script>

    <!-- TODO(scotttodd): url param inputs -->
    <!-- TODO(scotttodd): button for inputs -->
  </body>
</html>