#include <ctime>
#include <cstdlib>

#include <chrono>

#include "nnlayer3.inc.cpp"
#include "nnlayer3.mt.inc.cpp"


long long timeUs() {
  static std::chrono::steady_clock::time_point begin = std::chrono::steady_clock::now();
  return (long long)std::chrono::duration_cast<std::chrono::microseconds>( std::chrono::steady_clock::now() - begin ).count();
}


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.countNeurons() && l.back().countNeurons());

  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.countNeurons();
  int sizeY = l.back().countNeurons();
  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;

  long long timeStartUs = timeUs();
  for(int i = 0; i < blockCount; ++i) {
    long long timeBlockStartUs = timeUs();

    for(int j = 0; j < blockSize; ++j)
      shuffle[j] = rand()%count;
    double res = tmt.train(4);

    long long dt = timeUs() - timeBlockStartUs;

    printf("%4d, total %7d, avg.result %f, time: %f\n", i+1, (i+1)*blockSize, res, dt*0.000001);
    if ( ((i+1)%100) == 0 || i+1 == blockCount ) {
      if (!l.saveAll(outfile)) return delete[] data, delete[] shuffle, false;
      printf("  saved\n");
    }
  }

  long long dt = timeUs() - timeStartUs;
  printf("finished int time: %f\n", dt*0.000001);

  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, 28, 28,  1);
  new Layer(&l, 14, 14,  1, 28);
  new Layer(&l,  7,  7,  1, 14);
  new Layer(&l,  1,  1, 10,  7);

  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, 1000000, 0.01);

  return 0;
}