const len = 784;
const totalData = 1000;

const CAT = 0;
const RAINBOW = 3;
const TRAIN = 1;
const CAR = 4;
const ALARM = 2;
const EFIL = 5;

let catsData;
let trainsData;
let rainbowsData;
let carData;
let effilData;
let alarmData;

let cats = {};
let trains = {};
let rainbows = {};
let car = {};
let effil = {};
let alarm = {};
let nn;

function preload() {
  catsData = loadBytes('uploads/ozad2/cats1000.bin');
  trainsData = loadBytes('uploads/ozad2/trains1000.bin');
  //rainbowsData = loadBytes('/uploads/ozad2/rainbows1000.bin');
  //carData = loadBytes('data/car.bin');
  //effilData = loadBytes('uploads/ozad2/effil.bin');
  alarmData = loadBytes('uploads/ozad2/alarm.bin');
}


function setup() {
  createCanvas(500, 500);
	
  thehtml = "<button id='train'>train</button> <button id='test'>test</button> <button id='guess'>guess</button> <button id='clear'>clear</button><br>"+
  "<div>Draw either Alarm, Cat, Train.</div><span>Guess output is:=</span><div id = 'output'></div></br><div>Epoch:=<span id= 'epoch'></span></div><div>Test percentage:=<span id= 'testpercentage'></span></div>";
   AB.msg ( thehtml, 7 );
   
  background('gray');
    
   // console.log(catsData);
    //console.log(trainsData);
  // Preparing the data
  prepareData(cats, catsData, CAT);
  //prepareData(rainbows, rainbowsData, RAINBOW);
  prepareData(trains, trainsData, TRAIN);

  //prepareData(car, carData, CAR);
  
  //prepareData(effil, effilData, CAR);
  prepareData(alarm, alarmData, ALARM);


  // Making the neural network
  nn = new NeuralNetwork(784, 64, 3);

  // Randomizing the data
  let training = [];
  training = training.concat(cats.training);
  //training = training.concat(rainbows.training);
  training = training.concat(trains.training);
  //training = training.concat(effil.training);
  //training = training.concat(car.training);
  training = training.concat(alarm.training);

  let testing = [];
  testing = testing.concat(cats.testing);
  //testing = testing.concat(rainbows.testing);
  testing = testing.concat(trains.testing);
  //testing = testing.concat(car.testing);
  testing = testing.concat(alarm.testing);
  //testing = testing.concat(effil.testing);



let id = document.getElementById("train");
let epochCounter = 0;
$(id).click(function (event) {
    
    trainEpoch(training);
    epochCounter++;
    $("#epoch").text(epochCounter);
    console.log("Epoch: " + epochCounter);
    AB.removeLoading();
});

  /*let trainButton = select('#train');
  
  trainButton.mousePressed(function() {
   // console.log('training');
    trainEpoch(training);
    epochCounter++;
    console.log("Epoch: " + epochCounter);
  });*/

let id_test = document.getElementById("test");
$(id_test).click(function (event) {
    let percent = testAll(testing);
    console.log("Percent: " + nf(percent, 2, 2) + "%");
    $('#testpercentage').text(nf(percent, 2, 2)+"%");
});

/*  let testButton = select('#test');
  testButton.mousePressed(function() {
    
  });*/
  
  let id_guess = document.getElementById("guess");
$(id_guess).click(function (event) {
        let inputs = [];
    let img = get();
    img.resize(28, 28);
    img.loadPixels();
    for (let i = 0; i < len; i++) {
      let bright = img.pixels[i * 4];
      inputs[i] = (255 - bright) / 255.0;
    }

    let guess = nn.predict(inputs);
     console.log(guess);
    let m = max(guess);
    let classification = guess.indexOf(m);
    console.log(classification);
    if (classification === CAT) {
      console.log("cat");
      $("#output").text("Cat");
    } else if (classification === RAINBOW) {
      console.log("rainbow");
    } else if (classification === TRAIN) {
      console.log("train");
      $("#output").text("Train");
    }  else if (classification === CAR) {
      console.log("car");
    }  else if (classification === EFIL) {
      console.log("effil");
    } else if (classification === ALARM) {
      console.log("alarm");
      $("#output").text("Alarm Clock");
    }
});

  /*let guessButton = select('#guess');
  guessButton.mousePressed(function() {

    //image(img, 0, 0);
  });*/

  //let clearButton = select('#clear');
  
  let id_clear = document.getElementById("clear");
$(id_clear).click(function (event) {
    //background(255);
    background('gray');
});

  // for (let i = 1; i < 6; i++) {
  //   trainEpoch(training);
  //   console.log("Epoch: " + i);
  //   let percent = testAll(testing);
  //   console.log("% Correct: " + percent);
  // }
}


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


function trainEpoch(training) {
  shuffle(training, true);
  console.log(training);
  // Train for one epoch
  for (let i = 0; i < training.length; i++) {
    console.log(training[i]);
    //return;
    let data = training[i];
    console.log('data');
    console.log(data);
    let inputs = Array.from(data).map(x => x / 255);
    let label = training[i].label;
    let targets = [0, 0, 0];
    targets[label] = 1;
    // console.log(inputs);
    // console.log(targets);
    nn.train(inputs, targets);
  }
}

function testAll(testing) {

  let correct = 0;
  // Train for one epoch
  for (let i = 0; i < testing.length; i++) {
    // for (let i = 0; i < 1; i++) {
    let data = testing[i];
    let inputs = Array.from(data).map(x => x / 255);
    let label = testing[i].label;
    let guess = nn.predict(inputs);

    let m = max(guess);
    let classification = guess.indexOf(m);
    // console.log(guess);
    // console.log(classification);
    // console.log(label);

    if (classification === label) {
      correct++;
    }
  }
  let percent = 100 * correct / testing.length;
  return percent;

}



// Other techniques for learning

class ActivationFunction {
  constructor(func, dfunc) {
    this.func = func;
    this.dfunc = dfunc;
  }
}

let sigmoid = new ActivationFunction(
  x => 1 / (1 + Math.exp(-x)),
  y => y * (1 - y)
);

let tanh = new ActivationFunction(
  x => Math.tanh(x),
  y => 1 - (y * y)
);


class NeuralNetwork {
  /*
  * if first argument is a NeuralNetwork the constructor clones it
  * USAGE: cloned_nn = new NeuralNetwork(to_clone_nn);
  */
  constructor(in_nodes, hid_nodes, out_nodes) {
    if (in_nodes instanceof NeuralNetwork) {
      let a = in_nodes;
      this.input_nodes = a.input_nodes;
      this.hidden_nodes = a.hidden_nodes;
      this.output_nodes = a.output_nodes;

      this.weights_ih = a.weights_ih.copy();
      this.weights_ho = a.weights_ho.copy();

      this.bias_h = a.bias_h.copy();
      this.bias_o = a.bias_o.copy();
    } else {
      this.input_nodes = in_nodes;
      this.hidden_nodes = hid_nodes;
      this.output_nodes = out_nodes;

      this.weights_ih = new Matrix(this.hidden_nodes, this.input_nodes);
      this.weights_ho = new Matrix(this.output_nodes, this.hidden_nodes);
      this.weights_ih.randomize();
      this.weights_ho.randomize();

      this.bias_h = new Matrix(this.hidden_nodes, 1);
      this.bias_o = new Matrix(this.output_nodes, 1);
      this.bias_h.randomize();
      this.bias_o.randomize();
    }

    // TODO: copy these as well
    this.setLearningRate();
    this.setActivationFunction();

  }

  predict(input_array) {
    //console.log(input_array);
    // Generating the Hidden Outputs
    let inputs = Matrix.fromArray(input_array);
    //console.log(inputs);
    //console.log('weights');
   // console.log(this.weights_ih);
    let hidden = Matrix.multiply(this.weights_ih, inputs);
    
    //console.log('bias_h');
    //console.log(this.bias_h);
    hidden.add(this.bias_h);
    // activation function!
    hidden.map(this.activation_function.func);
    //console.log('hidden');
    //console.log(hidden);
    // Generating the output's output!
    let output = Matrix.multiply(this.weights_ho, hidden);
    //console.log('output');
    output.add(this.bias_o);
    output.map(this.activation_function.func);

    //console.log(output);

    // Sending back to the caller!
    return output.toArray();
  }

  setLearningRate(learning_rate = 0.01) {
    this.learning_rate = learning_rate;
  }

  setActivationFunction(func = sigmoid) {
    this.activation_function = func;
  }
  /*setActivationFunction(func = tanh) {
    this.activation_function = func;
  }*/

  train(input_array, target_array) {
    // Generating the Hidden Outputs
    //console.log(input_array);
    //console.log(target_array);
    let inputs = Matrix.fromArray(input_array);
    let hidden = Matrix.multiply(this.weights_ih, inputs);
    hidden.add(this.bias_h);
    // activation function!
    hidden.map(this.activation_function.func);

    // Generating the output's output!
    let outputs = Matrix.multiply(this.weights_ho, hidden);
    outputs.add(this.bias_o);
    outputs.map(this.activation_function.func);

    // Convert array to matrix object
    let targets = Matrix.fromArray(target_array);

    // Calculate the error
    // ERROR = TARGETS - OUTPUTS
    let output_errors = Matrix.subtract(targets, outputs);

    // let gradient = outputs * (1 - outputs);
    // Calculate gradient
    let gradients = Matrix.map(outputs, this.activation_function.dfunc);
    gradients.multiply(output_errors);
    gradients.multiply(this.learning_rate);


    // Calculate deltas
    let hidden_T = Matrix.transpose(hidden);
    let weight_ho_deltas = Matrix.multiply(gradients, hidden_T);

    // Adjust the weights by deltas
    this.weights_ho.add(weight_ho_deltas);
    // Adjust the bias by its deltas (which is just the gradients)
    this.bias_o.add(gradients);

    // Calculate the hidden layer errors
    let who_t = Matrix.transpose(this.weights_ho);
    let hidden_errors = Matrix.multiply(who_t, output_errors);

    // Calculate hidden gradient
    let hidden_gradient = Matrix.map(hidden, this.activation_function.dfunc);
    hidden_gradient.multiply(hidden_errors);
    hidden_gradient.multiply(this.learning_rate);

    // Calcuate input->hidden deltas
    let inputs_T = Matrix.transpose(inputs);
    let weight_ih_deltas = Matrix.multiply(hidden_gradient, inputs_T);

    this.weights_ih.add(weight_ih_deltas);
    // Adjust the bias by its deltas (which is just the gradients)
    this.bias_h.add(hidden_gradient);
  }

  serialize() {
    return JSON.stringify(this);
  }

  static deserialize(data) {
    if (typeof data == 'string') {
      data = JSON.parse(data);
    }
    let nn = new NeuralNetwork(data.input_nodes, data.hidden_nodes, data.output_nodes);
    nn.weights_ih = Matrix.deserialize(data.weights_ih);
    nn.weights_ho = Matrix.deserialize(data.weights_ho);
    nn.bias_h = Matrix.deserialize(data.bias_h);
    nn.bias_o = Matrix.deserialize(data.bias_o);
    nn.learning_rate = data.learning_rate;
    return nn;
  }


  // Adding function for neuro-evolution
  copy() {
    return new NeuralNetwork(this);
  }

  // Accept an arbitrary function for mutation
  mutate(func) {
    this.weights_ih.map(func);
    this.weights_ho.map(func);
    this.bias_h.map(func);
    this.bias_o.map(func);
  }



}


// let m = new Matrix(3,2);


class Matrix {
  constructor(rows, cols) {
    this.rows = rows;
    this.cols = cols;
    this.data = Array(this.rows).fill().map(() => Array(this.cols).fill(0));
  }

  copy() {
    let m = new Matrix(this.rows, this.cols);
    for (let i = 0; i < this.rows; i++) {
      for (let j = 0; j < this.cols; j++) {
        m.data[i][j] = this.data[i][j];
      }
    }
    return m;
  }

  static fromArray(arr) {
    return new Matrix(arr.length, 1).map((e, i) => arr[i]);
  }

  static subtract(a, b) {
    if (a.rows !== b.rows || a.cols !== b.cols) {
      console.log('Columns and Rows of A must match Columns and Rows of B.');
      return;
    }

    // Return a new Matrix a-b
    return new Matrix(a.rows, a.cols)
      .map((_, i, j) => a.data[i][j] - b.data[i][j]);
  }

  toArray() {
    let arr = [];
    for (let i = 0; i < this.rows; i++) {
      for (let j = 0; j < this.cols; j++) {
        arr.push(this.data[i][j]);
      }
    }
    return arr;
  }

  randomize() {
    return this.map(e => Math.random() * 2 - 1);
  }

  add(n) {
    if (n instanceof Matrix) {
      if (this.rows !== n.rows || this.cols !== n.cols) {
        console.log('Columns and Rows of A must match Columns and Rows of B.');
        return;
      }
      return this.map((e, i, j) => e + n.data[i][j]);
    } else {
      return this.map(e => e + n);
    }
  }

  static transpose(matrix) {
    return new Matrix(matrix.cols, matrix.rows)
      .map((_, i, j) => matrix.data[j][i]);
  }

  static multiply(a, b) {
    //console.log('static multiply');
    console.log(a);
    console.log(b);
    // Matrix product
    if (a.cols !== b.rows) {
      console.log('Columns of A must match rows of B.');
      return;
    }

    return new Matrix(a.rows, b.cols)
      .map((e, i, j) => {
        // Dot product of values in col
        let sum = 0;
        for (let k = 0; k < a.cols; k++) {
          sum += a.data[i][k] * b.data[k][j];
        }
        return sum;
      });
  }

  multiply(n) {
    //console.log('multiply');
    if (n instanceof Matrix) {
      if (this.rows !== n.rows || this.cols !== n.cols) {
        console.log('Columns and Rows of A must match Columns and Rows of B.');
        return;
      }

      // hadamard product
      return this.map((e, i, j) => e * n.data[i][j]);
    } else {
      // Scalar product
      return this.map(e => e * n);
    }
  }

  map(func) {
    // Apply a function to every element of matrix
    for (let i = 0; i < this.rows; i++) {
      for (let j = 0; j < this.cols; j++) {
        let val = this.data[i][j];
        this.data[i][j] = func(val, i, j);
      }
    }
    return this;
  }

  static map(matrix, func) {
    // Apply a function to every element of matrix
    return new Matrix(matrix.rows, matrix.cols)
      .map((e, i, j) => func(matrix.data[i][j], i, j));
  }

  print() {
    console.table(this.data);
    return this;
  }

  serialize() {
    return JSON.stringify(this);
  }

  static deserialize(data) {
    if (typeof data == 'string') {
      data = JSON.parse(data);
    }
    let matrix = new Matrix(data.rows, data.cols);
    matrix.data = data.data;
    return matrix;
  }
}

if (typeof module !== 'undefined') {
  module.exports = Matrix;
}

function loadByte(file, callback)
{
    var self = this;
  var data = {};
  var oReq = new XMLHttpRequest();
  oReq.open("GET", file, true);
  oReq.responseType = "arraybuffer";
  oReq.onload = function(oEvent) {
    var arrayBuffer = oReq.response;
    if (arrayBuffer) {
      data.bytes = new Uint8Array(arrayBuffer);
      if (callback) {
        callback(data);
      }
      self._decrementPreload();
    }
  }
  oReq.send(null);
 // console.log(data);
  return data;
}

function prepareData(category, data, label) {
  category.training = [];
  category.testing = [];
  for (let i = 0; i < totalData; i++) {
    let offset = i * len;
    let threshold = floor(0.8 * totalData);
    if (i < threshold) {
      category.training[i] = data.bytes.subarray(offset, offset + len);
      category.training[i].label = label;
    } else {
      category.testing[i - threshold] = data.bytes.subarray(offset, offset + len);
      category.testing[i - threshold].label = label;
    }
  }
  //console.log(category);
}