#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;

}