#include <ctime></ctime>

#include <cstdlib></cstdlib>

#include <cstdio></cstdio>

#include <algorithm></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;

}