#include <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";
  (new LayerSimple<funcSigmoidExp>    ( *nl, Layout( 64)      ))->filename = FILENAME "2";
  (new LayerSimple<funcSigmoidExp>    ( *nl, Layout( 10)      ))->filename = FILENAME "3";
  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;
}