#include <ctime></ctime>

#include <cstdlib></cstdlib>

#include "nnlayer2.inc.cpp"

#include "nnlayer2.conv.inc.cpp"

#include "nnlayer2.mt.inc.cpp"

bool train(const char *infile, const char *outfile, Layer &l, int blockSize, int totalCount, Real trainRatio) {

  assert(blockSize > 0);

  int blockCount = totalCount/blockSize;

  assert(blockCount > 0);

  assert(!l.prev);

  assert(l.size && l.back().size);

  printf("load training data\n");

  FILE *f = fopen(infile, "rb");

  if (!f)

    return printf("cannot open file '%s' for read\n", infile), false;

  fseek(f, 0, SEEK_END);

  int fs = ftell(f);

  fseek(f, 0, SEEK_SET);

  int sizeX = l.size;

  int sizeY = l.back().size;

  int count = fs/(sizeX+1);

  if (count < blockSize)

    return printf("file '%s' is lesser minimal size\n", infile), fclose(f), false;

  unsigned char *data = new unsigned char[(sizeX + sizeY)*count];

  memset(data, 0, (sizeX + sizeY)*count);

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

    unsigned char *d = data + (sizeX + sizeY)*i;

    if (!fread(d, sizeX+1, 1, f) || d[sizeX] >= sizeY)

      return printf("cannot read from file '%s'\n", infile), delete[] data, fclose(f), false;

    d += sizeX;

    unsigned char c = *d;

    *d = 0;

    d[c] = 255;

  }

  fclose(f);

  printf("train %d x %d = %d, ratio: %f\n", blockCount, blockSize, blockCount*blockSize, trainRatio);

  int *shuffle = new int[blockSize];

  TrainMT tmt;

  tmt.layer = &l;

  tmt.dataX = data;

  tmt.dataY = data + sizeX;

  tmt.strideX = tmt.strideY = sizeX + sizeY;

  tmt.shuffle = shuffle;

  tmt.count = blockSize;

  tmt.trainRatio = trainRatio;

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

    for(int j = 0; j < blockSize; ++j)

      shuffle[j] = rand()%count;

    Real res = tmt.train(8);

    printf("%4d, total %7d, avg.result %f\n", i+1, (i+1)*blockSize, res);

    if ( ((i+1)%10) == 0 || i+1 == blockCount ) {

      if (!l.saveAll(outfile)) return delete[] data, delete[] shuffle, false;

      printf("  saved\n");

    }

  }

  delete[] shuffle;

  delete[] data;

  return true;

}

int main() {

  srand(time(NULL));

  const char *infile = "data/symbols-data.bin"; // 28x28

  const char *outfile = "data/output/weights.bin";

  printf("create neural network\n");

  Layer l(nullptr, 784);

  createConv(l, 28, 28, 1, 22, 22, 1,  60);

  createConv(l, 22, 22, 1, 14, 14, 1, 100);

  createConv(l, 14, 14, 1,  4,  4, 1, 140);

  createConv(l, 4, 4, 1, 1, 1, 10, 16);

  printf("  neurons: %d, links %d, memSize: %llu\n", l.totalNeurons(), l.totalLinks(), (unsigned long long)l.totalMemSize());

  //printf("try load previously saved network\n");

  //l.loadAll(outfile);

  printf("train\n");

  train(infile, outfile, l, 10000, 2000000, 0.1);

  return 0;

}