#ifndef TRAIN_IMAGE_INC_CPP

#define TRAIN_IMAGE_INC_CPP

#include "train.inc.cpp"

#include "layer.simple.inc.cpp"

class TrainerImage: public Trainer {

protected:

  std::vector<unsigned char=""> data;</unsigned>

  std::vector<unsigned int=""> shuffle;</unsigned>

  const char *datafile;

  const char *outfile;

  Layout ofl, obl;

  Layout::List oflist, oblist;

  int stride, count;

public:

  TrainerImage(): stride(), count() { }

  bool configure(const char *datafile, const char *outfile) {

    this->datafile = datafile;

    this->outfile = outfile;

  }

    data.clear();

    FILE *f = fopen(filename, "rb");

    if (!f)

      return printf("cannot open file for read: %s\n", filename), false;

    fseek(f, 0, SEEK_END);

    size_t fs = ftello(f);

    fseek(f, 0, SEEK_SET);

    data.resize(fs, 0);

    if (!fread(data.data(), fs, 1, f))

      return printf("cannot read from file: %s\n", filename), fclose(f), data.clear(), false;

    fclose(f);

    return true;

  }

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");

}

protected:

  bool prepare() override {

    ofl = optimizeLayoutSimple(fl->layout);

    obl = optimizeLayoutSimple(bl->layout);

    assert(ofl && obl);

    assert(ofl.getActiveCount() == obl.getActiveCount());

    ofl.split(oflist, threadsCount);

    obl.split(oblist, threadsCount);

    stride = ofl.getActiveCount() + 1;

    count = data.size()/stride;

    if (count <= 0) return false;

    shuffle.resize(count);

    for(int i = 0; i < count; ++i)

      shuffle[i] = i;

    return true;

  }

  bool prepareBlock() override {

    int cnt = itersPerBlock > count ? count : itersPerBlock;

    for(int i = 0; i < cnt; ++i) {

      int j = rand()%count;

      if (i != j) std::swap(shuffle[i], shuffle[j]);

    }

    return true;

  }

  void loadData(Barrier &barrier, int, int iter) override {

    struct I: public Iter {

      typedef const unsigned char* Type;

      static inline void iter4(Neuron &n, Type d, AccumType&) { n.v = *d/(NeuronReal)255; }

    };

    const unsigned char *id = data.data() + shuffle[iter%count]*stride;

    iterateNeurons2(oflist[barrier.tid], ofl, fl->neurons, id);

  }

  AccumReal verifyDataMain(int, int iter) override {

    struct I: public Iter {

      typedef int Type;

      struct AccumType { int ri, mi; NeuronReal m; };

      static inline void iter4(Neuron &n, Type d, AccumType &a) {

        NeuronReal v1 = d == a.ri;

        NeuronReal v0 = n.v;

        n.d *= v1 - v0;

        if (a.m < v0) { a.m = v0; a.mi = d; }

      }

    };

    I::AccumType a = { data[ (shuffle[iter%count] + 1)*stride - 1 ], 0, 0 };

    iterateNeurons2(obl, obl, bl->neurons, 0, 1, &a);

    return a.mi != a.ri;

  }

};

#endif