var MODEL_LOADING = 0;
var MODEL_READY = 1;
var MODEL_LOAD_ERROR = 2;
var MODEL_WORKING = 3;
var modelStatus = MODEL_LOADING;
var modelStatusColors = [[200, 200, 0], [32, 128, 32], [200, 0, 0], [32, 32, 255]];
var modelOutput = null;
var imgMask = null;
var curModelImage = null;
var running = false;

// we start with some probes. probeQueue = null when done
var probeQueue = [];
var probeCoords = null;

// http://stackoverflow.com/a/8843728/1010653
function norm_ppf(p) {
    var a1 = -39.6968302866538, a2 = 220.946098424521, a3 = -275.928510446969;
    var a4 = 138.357751867269, a5 = -30.6647980661472, a6 = 2.50662827745924;
    var b1 = -54.4760987982241, b2 = 161.585836858041, b3 = -155.698979859887;
    var b4 = 66.8013118877197, b5 = -13.2806815528857, c1 = -7.78489400243029E-03;
    var c2 = -0.322396458041136, c3 = -2.40075827716184, c4 = -2.54973253934373;
    var c5 = 4.37466414146497, c6 = 2.93816398269878, d1 = 7.78469570904146E-03;
    var d2 = 0.32246712907004, d3 = 2.445134137143, d4 = 3.75440866190742;
    var p_low = 0.02425, p_high = 1 - p_low;
    var q, r;
    var retVal;

    if ((p < 0) || (p > 1))
    {
        alert("NormSInv: Argument out of range.");
        retVal = 0;
    }
    else if (p < p_low)
    {
        q = Math.sqrt(-2 * Math.log(p));
        retVal = (((((c1 * q + c2) * q + c3) * q + c4) * q + c5) * q + c6) / ((((d1 * q + d2) * q + d3) * q + d4) * q + 1);
    }
    else if (p <= p_high)
    {
        q = p - 0.5;
        r = q * q;
        retVal = (((((a1 * r + a2) * r + a3) * r + a4) * r + a5) * r + a6) * q / (((((b1 * r + b2) * r + b3) * r + b4) * r + b5) * r + 1);
    }
    else
    {
        q = Math.sqrt(-2 * Math.log(1 - p));
        retVal = -(((((c1 * q + c2) * q + c3) * q + c4) * q + c5) * q + c6) / ((((d1 * q + d2) * q + d3) * q + d4) * q + 1);
    }

    return retVal;
}

// in browser, URLs can be relative or absolute
const model = new KerasJS.Model({
  filepaths: {
    model: 'mnist_03_decoder.json',
    weights: 'mnist_03_decoder_weights_weights.buf',
    metadata: 'mnist_03_decoder_weights_metadata.json'
  },
  gpu: true
})

var origImgMask;
function preload() {
  origImgMask = loadImage("mask56.png");
}

function setup () {
  imgMask = createImage(56, 56);
  imgMask.loadPixels();
  origImgMask.loadPixels();
  // print(imgMask.pixels.length + " and " + origImgMask.pixels.length)
  for (var i = 3; i < imgMask.pixels.length; i+=4) {
    imgMask.pixels[i] = 255-origImgMask.pixels[i];
  }
  imgMask.updatePixels();

  createCanvas(960, 500);
  background(32);
  curModelImage = createImage(56, 56);
  model.ready()
    .then(() => {
      // print("Model says it is ready")
      modelStatus = MODEL_READY;
      // redraw();
    })
    .catch(err => {
      // handle error
      modelStatus = MODEL_LOAD_ERROR;
      // redraw();
    })
  probeQueue.push([(0+0.5)*width/10, (0+0.5)*height/5]);
  for(var i=0;i<10;i++) {
    for(var j=0;j<5;j++) {
      probeQueue.push([(i+0.5)*width/10, (j+0.5)*height/5]);
    }
  }
}

function mouseMoved() {
  if(running) {
    redraw();
  }
}

function mouseDragged() {
  if(running) {
    probeHere(mouseX, mouseY);
  }
}

function mouseClicked() {
  if(running) {
    probeHere(mouseX, mouseY);
  }
}

function canvasPos2LatentSpace(x, y) {
    var coordX = map(x, 0, width, 0.05, 0.95);
    var coordY = map(y, 0, height, 0.05, 0.95);
    return [norm_ppf(coordX), norm_ppf(coordY)];
}

function probeHere(x, y) {
  print("MOD STATUS: " + modelStatus);
  if(modelStatus == MODEL_READY) {
    modelStatus = MODEL_WORKING;
    var coords = canvasPos2LatentSpace(x, y);
    const inputData = { input_1: new Float32Array(coords) };
    model.predict(inputData).then(outputData => {
      modelOutput = outputData['convolution2d_5'];
      modelStatus = MODEL_READY;
      redraw();
      print("PREDICTED");
    });
  }
}

function draw () {
  if(modelStatus == MODEL_READY && modelOutput == null &&
      probeQueue != null && probeQueue.length > 0) {
    probeCoords = probeQueue.shift();
    probeHere(probeCoords[0], probeCoords[1]);
  }

  fill(modelStatusColors[modelStatus]);
  ellipse(20, height-20, 10, 10);

  if(modelOutput != null) {
    curModelImage.loadPixels();
    for (i = 0; i < 27; i++) {
      for (j = 0; j < 27; j++) {
        var val = (255 * modelOutput[j*27 + i]);
        curModelImage.set(i*2, j*2, color(val));
        curModelImage.set(i*2+1, j*2, color(val));
        curModelImage.set(i*2, j*2+1, color(val));
        curModelImage.set(i*2+1, j*2+1, color(val));
      }
    }
    curModelImage.updatePixels();
    curModelImage.mask(imgMask);
    if (probeCoords != null) {
      image(curModelImage, probeCoords[0]-28, probeCoords[1]-28);      
    }
    else if (mouseX != 0 && mouseY != 0) {
      image(curModelImage, mouseX-28, mouseY-28);
    }
    modelOutput = null;
    if(!running && probeQueue.length == 0) {
      probeQueue = null;
      probeCoords = null;
      running = true;
      noLoop();
    }
  }
}

function keyTyped() {
  if (key == '!') {
    saveBlocksImages();    
  }
  else if (key == '@') {
    saveBlocksImages(true);    
  }
}