#include <ctime>
#include <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;
}