#include <math.h>

#include <stdio.h>

#include <stdlib.h>

#include <helianthus.h>

#include "nnlayer.inc.c"

NeuralLayer *nl;

Framebuffer fb, fbMin;

Animation fbAnim, fbMinAnim;

int wasPressed;

double prevX, prevY;

void prepareImage() {

  int w, h;

  unsigned char *pixels = NULL;

  saveState();

  target(fb);

  imageFromViewport(&w, &h, &pixels);

  restoreState();

  if (!pixels) return;

  int x0 = w, y0 = h, x1 = 0, y1 = 0;

  for(int y = 0; y < h; ++y) {

    for(int x = 0; x < w; ++x) {

      if (imageGetPixel(w, h, pixels, x, y) != 0x000000ff) {

        if (x0 > x) x0 = x;

        if (x1 < x) x1 = x;

        if (y0 > y) y0 = y;

        if (y1 < y) y1 = y;

      }

    }

  }

  free(pixels);

  pixels = NULL;

  if (x1 < x0 || y1 < y0) return;

  double wx = x1 - x0 + 1;

  double wy = y1 - y0 + 1;

  double s = 20.0/(wx > wy ? wx : wy);

  double cx = (x0 + x1)/2.0;

  double cy = (y0 + y1)/2.0;

  double xx = 14 - s*cx;

  double yy = 14 - s*cy;

  double ww = s*w;

  double hh = s*h;

  saveState();

  target(fbMin);

  noStroke();

  rectTextured(fbAnim, xx, yy, ww, hh);

  imageFromViewport(&w, &h, &pixels);

  restoreState();

  if (!pixels) return;

  double *p = nl->a;

  for(int y = 0; y < h; ++y)

    for(int x = 0; x < w; ++x)

      *p++ = colorGetValue(imageGetPixel(w, h, pixels, x, y));

}

void init() {

  background(COLOR_BLACK);

  stroke(COLOR_WHITE);

  fb = createFramebufferEx(512, 512, NULL, FALSE, FALSE, TRUE);

  fbMin = createFramebufferEx(28, 28, NULL, FALSE, FALSE, TRUE);

  fbAnim = createAnimationFromFramebuffer(fb);

  fbMinAnim = createAnimationFromFramebuffer(fbMin);

  saveState();

  target(fb);

  clear();

  target(fbMin);

  clear();

  restoreState();

  nl = nlNew(NULL, 784, 0);

  nlNew(nl, 256, 0);

  nlNew(nl, 128, 0);

  nlNew(nl, 64, 0);

  nlNew(nl, 64, 0);

  nlNew(nl, 10, 0);

  nlLoad(nl, "data/weights.bin");

}

void draw() {

  saveState();

  if (mouseDown("left")) {

    double x = mouseX(), y = mouseY();

    if (!wasPressed) prevX = x, prevY = y;

    saveState();

    strokeWidth(32);

    target(fb);

    line(prevX, prevY, x, y);

    restoreState();

    prevX = x, prevY = y;

    wasPressed = TRUE;

  } else {

    wasPressed = FALSE;

  }

  if (keyWentDown("space")) {

    prepareImage();

    nlPass(nl);

    saveState();

    target(fb);

    clear();

    restoreState();

  }

  noStroke();

  rectTextured(fbAnim, 0, 0, 512, 512);

  stroke(COLOR_WHITE);

  rectTextured(fbMinAnim, 16, 16, 28, 28);

  noFill();

  NeuralLayer *nlb = nlBack(nl);

  textSize(8);

  int res = 0;

  for(int i = 0; i < 10; ++i) {

    if (nlb->a[i] > nlb->a[res]) res = i;

    textf(16, 90+8*i, "%d: %lf", i, nlb->a[i]);

  }

  textSize(16);

  textf(16, 60, "%d", res);

  restoreState();

}

int main(int largc, char **largv) {

  windowSetVariableFrameRate();

  windowSetInit(&init);

  windowSetDraw(&draw);

  windowRun();

  return 0;

}