// Shape Classifier (Mouse)
// Coding Challenge
// The Coding Train / Daniel Shiffman
// https://thecodingtrain.com/CodingChallenges/158-shape-classifier.html
// https://youtu.be/3MqJzMvHE3E

// Generate Dataset: https://github.com/CodingTrain/website/tree/gh-pages/CodingChallenges/CC_158_Shape_Classifier/dataset
// Generate Dataset (port): https://editor.p5js.org/codingtrain/sketches/7leVIzy5l
// Training: https://github.com/CodingTrain/website/tree/gh-pages/CodingChallenges/CC_158_Shape_Classifier/training
// Mouse: https://editor.p5js.org/codingtrain/sketches/JgLVfCS4E
// Webcam: https://editor.p5js.org/codingtrain/sketches/2hZGBkqqq



var nodeDoctype = document.implementation.createDocumentType(
 'html',
 "script src=https://cdn.jsdelivr.net/npm/p5@1.4.0/lib/p5.min.js",
 "script src=https://unpkg.com/ml5@0.6.0/dist/ml5.min.js",
 "meta charset=utf-8",
 "script src=sketch.js"
);


let circles = [];
let squares = [];
let triangles = [];

let shapeClassifier;
let canvas;
let resultsDiv;
let inputImage;
let clearButton;

function preload() {
  for (let i = 0; i < 4; i++) {
    let index = nf(i + 1, 4, 0);
    circles[i] = loadImage(`/uploads/vrushali/circle${index}.png`);
    squares[i] = loadImage(`/uploads/vrushali/square${index}.png`);
    triangles[i] = loadImage(`/uploads/vrushali/triangle${index}.png`);
  }
}


function setup() {
  canvas = createCanvas(400, 400);
  let options = {
    inputs: [64, 64, 4],
    task: 'imageClassification',
    debug: true
  };
  
  
  $.getScript ( "/uploads/vrushali/index.html", function()
 {
  
            console.log ("All JS loaded");
        
 });
  
  shapeClassifier = ml5.neuralNetwork(options);
  
  for (let i = 0; i < circles.length; i++) {
    shapeClassifier.addData({ image: circles[i] }, { label: 'circle' });
    shapeClassifier.addData({ image: squares[i] }, { label: 'square' });
    shapeClassifier.addData({ image: triangles[i] }, { label: 'triangle' });
  }
  shapeClassifier.normalizeData();
  shapeClassifier.train({ epochs: 50 }, finishedTraining);
  
  
  const modelDetails = {
    model: '/uploads/vrushali/model.json',
    metadata: '/uploads/vrushali/model_meta.json',
    weights: '/uploads/vrushali/model.weights.bin'
  };
  background(255);
  clearButton = createButton('clear');
  clearButton.mousePressed(function() {
    background(255);
  });
  resultsDiv = createDiv('loading model');
  inputImage = createGraphics(64, 64);
  shapeClassifier.load(modelDetails, modelLoaded);
}

function modelLoaded() {
  console.log('model ready!');
  classifyImage();
}

function classifyImage() {
  inputImage.copy(canvas, 0, 0, 400, 400, 0, 0, 64, 64);
  //image(inputImage, 0, 0);
  shapeClassifier.classify(
    {
      image: inputImage
    },
    gotResults
  );
}

function gotResults(err, results) {
  if (err) {
    console.error(err);
    return;
  }

  let label = results[0].label;
  let confidence = nf(100 * results[0].confidence, 2, 0);

  resultsDiv.html(`${label} ${confidence}%`);

  //console.log(results);
  classifyImage();
}


function finishedTraining() {
  console.log('finished training!');
  shapeClassifier.save();
}


function draw() {
  if (mouseIsPressed) {
    strokeWeight(8);
    line(mouseX, mouseY, pmouseX, pmouseY);
  }

  // stroke(0);
  // noFill();
  // strokeWeight(4);
  // rectMode(CENTER);
  // rect(width/2, height/2, 40);
}