#include <ctime>
#include <cstdlib>
#include <cstdio>

#include <algorithm>

#include "nnlayer.lnk.inc.cpp"
#include "tga.inc.cpp"


void imgTrain(Layer &l, const char *datafile, int size, const char *outfile, double trainRatio, int count) {
  Layer &bl = l.back();

  assert(!l.prev);
  assert(datafile);
  assert(count > 0 && size > 0);
  assert(l.size == size);
  assert(bl.size == size);

  int blockSize = 1000;//1024*1024*1024/size;
  assert(blockSize > 0);

  FILE *f = fopen(datafile, "rb");
  if (!f)
    { printf("cannot open file: %s\n", datafile); return; }
  fseeko64(f, 0, SEEK_END);
  long long fsize = ftello64(f);
  int xCount = (int)(fsize/size);
  if (xCount <= 0)
    { printf("no tests in file: %s\n", datafile); return; }

  int *block = new int[blockSize*2];
  int *shuffle = block + blockSize;
  double *results = new double[blockSize];
  unsigned char *blockData = new unsigned char[(blockSize + 1)*size];
  unsigned char *blockResData = blockData + blockSize*size;
  bool err = false;

  for(int j = 0; j < blockSize; ++j)
    { shuffle[j] = j; results[j] = 0; }

  int blocksCount = (count - 1)/blockSize + 1;

  printf("training %d (%d x %d blocks), tests: %d, ratio: %f:\n", blocksCount*blockSize, blocksCount, blockSize, xCount, trainRatio);

  double avgSum = 0;
  for(int i = 0; i < blocksCount; ++i) {
    for(int j = 0; j < blockSize; ++j) {
      block[j] = rand()%xCount;
      std::swap(shuffle[i], shuffle[rand()%blockSize]);
    }
    std::sort(block, block + blockSize);

    for(int j = 0; j < blockSize; ++j) {
      fseeko64(f, block[j]*(long long)size, SEEK_SET);
      if (!fread(blockData + j*size, size, 1, f))
        { printf("cannot read data from file: %s\n", datafile); err = true; break; }
    }
    if (err) break;

    printf("  next data block loaded\n");

    double sumQ = 0;
    for(int j = 0; j < blockSize; ++j) {
      unsigned char *data = blockData + shuffle[j]*size;
      for(double *ia = l.a, *e = ia + l.size; ia < e; ++ia, ++data)
        *ia = *data/255.0;

      double firstQ = 0, q = 0;
      for(int repeat = 0; repeat < 1; ++repeat) {
        l.pass();

        for(double *ia = l.a, *iba = bl.a, *ibda = bl.da, *e = ia + l.size; ia < e; ++ia, ++iba, ++ibda) {
          double d = *ia - *iba;
          *ibda = d;
          q += d*d;
        }
        q /= size;
        if (!repeat) firstQ = q;

        bl.backpass(trainRatio);
      }

      sumQ += firstQ;
      avgSum += firstQ - results[j];
      results[j] = firstQ;
      int avgCnt = i ? blockSize : j + 1;
      printf("  %4d: total: %6d, avg result: %f, last result: %f -> %f\n", j+1, i*blockSize+j+1, avgSum/avgCnt, firstQ, q);
    }

    printf("%4d: total: %6d, avg result: %f\n", i+1, (i+1)*blockSize, sumQ/blockSize);

    if (outfile && !l.save(outfile))
      { printf("cannot save neural network weights to file: %s\n", outfile); err = true; break; }

    unsigned char *data = blockResData;
    for(double *iba = bl.a, *e = iba + bl.size; iba < e; ++iba, ++data)
      *data = (unsigned char)(*iba*255.999);
    tgaSave("data/output/sampleX.tga", blockData + shuffle[blockSize-1]*size, 256, 256, 3);
    tgaSave("data/output/sampleY.tga", blockResData, 256, 256, 3);
  }

  delete[] block;
  delete[] results;
  delete[] blockData;

  printf("finished\n");
}


int main() {
  srand(time(NULL));

  //const char *datafile = "data/img512-data.bin";
  //const char *outfile = "data/output/img512-weights.bin";
  const char *datafile = "data/img256-data.bin";
  const char *outfile = "data/output/img256-weights.bin";

  printf("create neural network\n");
  //Layer l(nullptr, 512*512*3);
  //new LayerLinkConvolution(l, 512, 512, 3, 256, 256, 3,  32);
  //new LayerLinkConvolution(l, 256, 256, 3, 128, 128, 3,  64);
  //new LayerLinkConvolution(l, 128, 128, 3,  64,  64, 3, 128);
  //new LayerLinkConvolution(l,  64,  64, 3,  32,  32, 3, 256);
  //new LayerLinkConvolution(l,  32,  32, 3,  16,  16, 4, 256);
  //new LayerLinkConvolution(l,  16,  16, 4,  16,  16, 4, 256);
  //new LayerLinkConvolution(l,  16,  16, 4,  16,  16, 4, 256);
  //new LayerLinkConvolution(l,  16,  16, 4,  32,  32, 3, 256);
  //new LayerLinkConvolution(l,  32,  32, 3,  64,  64, 3, 128);
  //new LayerLinkConvolution(l,  64,  64, 3, 128, 128, 3,  64);
  //new LayerLinkConvolution(l, 128, 128, 3, 256, 256, 3,  32);
  //new LayerLinkConvolution(l, 256, 256, 3, 512, 512, 3,  16);

  Layer l(nullptr, 256*256*3);
  new LayerLinkConvolution(l, 256, 256, 3, 128, 128, 3,  32);
  new LayerLinkConvolution(l, 128, 128, 3,  64,  64, 3,  64);
  new LayerLinkConvolution(l,  64,  64, 3,  32,  32, 3, 128);
  new LayerLinkConvolution(l,  32,  32, 3,  16,  16, 4, 256);
  new LayerLinkConvolution(l,  16,  16, 4,  16,  16, 4, 256);
  new LayerLinkConvolution(l,  16,  16, 4,  16,  16, 4, 256);
  new LayerLinkConvolution(l,  16,  16, 4,  32,  32, 3, 128);
  new LayerLinkConvolution(l,  32,  32, 3,  64,  64, 3,  64);
  new LayerLinkConvolution(l,  64,  64, 3, 128, 128, 3,  32);
  new LayerLinkConvolution(l, 128, 128, 3, 256, 256, 3,  16);

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

  printf("try load previously saved network\n");
  l.load(outfile);

  printf("train\n");
  imgTrain(l, datafile, l.size, outfile, 0.1, 1000000);

  return 0;
}