#include <helianthus.h></helianthus.h>

#include "layer.all.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;

  Neuron *in = nl->neurons;

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

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

      (in++)->v = 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();

  #define FILENAME "data/weights-digit.bin"

  nl = new Layer(                   nullptr, Layout(28, 28));

  (new LayerSimple<funcsigmoidexp>    ( *nl, Layout(256)      ))->filename = FILENAME "1";</funcsigmoidexp>

  (new LayerSimple<funcsigmoidexp>    ( *nl, Layout( 64)      ))->filename = FILENAME "2";</funcsigmoidexp>

  (new LayerSimple<funcsigmoidexp>    ( *nl, Layout( 10)      ))->filename = FILENAME "3";</funcsigmoidexp>

  nl->load();

}

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();

    for(Layer *l = nl->next; l; l = l->next)

      l->testPass();

    saveState();

    target(fb);

    clear();

    restoreState();

  }

  noStroke();

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

  stroke(COLOR_WHITE);

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

  noFill();

  Layer &nlb = nl->back();

  textSize(8);

  int res = 0;

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

    if (nlb.neurons[i].v > nlb.neurons[res].v) res = i;

    textf(16, 90+8*i, "%d: %lf", i, nlb.neurons[i].v);

  }

  textSize(16);

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

  restoreState();

}

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

  windowSetVariableFrameRate();

  windowSetInit(&init);

  windowSetDraw(&draw);

  windowRun();

  return 0;

}