#include <math.h>
#include <stdio.h>
#include <stdlib.h>

#include <helianthus.h>

#include "nnlayer.inc.cpp"


Layer *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;

  int fw = framebufferGetWidth(fbMin);
  int fh = framebufferGetHeight(fbMin);

  double wx = x1 - x0 + 1;
  double wy = y1 - y0 + 1;
  double s = (fw - 4)/(double)(wx > wy ? wx : wy);
  double cx = (x0 + x1)/2.0;
  double cy = (y0 + y1)/2.0;

  double xx = fw/2 - s*cx;
  double yy = fh/2 - 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 = new Layer(nullptr, 784);
  new LayerAdd(*nl, 784);
  new LayerSimple(*nl, 256);
  new LayerSimple(*nl, 128);
  new LayerSimple(*nl, 64);
  new LayerSimple(*nl, 64);
  new LayerSimple(*nl, 10);

  //nl = new Layer(nullptr, 22*22);
  //new LayerConvolution(*nl, 10, 10, 4);
  //new LayerConvolution(*nl, 4, 4, 16);
  //new LayerConvolution(*nl, 1, 1, 64);
  //new LayerPlain(*nl, 64);
  //new LayerPlain(*nl, 32);
  //new LayerPlain(*nl, 10);

  nl->load("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();
    nl->pass();
    saveState();
    target(fb);
    clear();
    restoreState();
  }

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

  stroke(COLOR_WHITE);
  rectTextured(fbMinAnim, 16, 16, 22, 22);

  noFill();

  Layer &nlb = nl->back();
  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;
}