#ifndef TRAIN_CX4_INC_CPP
#define TRAIN_CX4_INC_CPP


#include "train.segment.inc.cpp"
#include "segment.cx4.inc.cpp"
#include "layer.inc.cpp"


class TrainerCx4: public TrainerSegment {
protected:
  FILE *f;
  std::vector<unsigned char> data;
  std::vector<NeuronReal> values;
  std::vector<NeuronReal> valuesMeasure;
  std::vector<unsigned char> tmpdata;
  std::vector<int> shuffle;
  
  Layout trainLayout;
  Layout measureLayout;
  
  size_t imageSize;
  size_t preparedImageSize;
  int imagesInFile;
  int imagesInMemory;
  
  volatile unsigned int seed;

public:
  Layer *layerFull;
  Layer *layerPre;
  int loadImagesCount;
  int blocksPerLoading;
  
  const char *infile;
  const char *cachefile;
  const char *outfile;
  
  TrainerCx4():
    f(),
    imageSize(),
    preparedImageSize(),
    imagesInFile(),
    imagesInMemory(),
    seed(),
    layerFull(),
    layerPre(),
    loadImagesCount(),
    blocksPerLoading(1),
    infile(),
    cachefile(),
    outfile() { }

protected:
  void preprocess(unsigned char *src, NeuronReal *dst) {
    struct IL: public Iter {
      typedef const unsigned char* DataType;
      static inline void iter4(Neuron &n, DataType d, DataAccumType&) { n.v = *d/(NeuronReal)255; }
    };
    struct IS: public Iter {
      typedef NeuronReal* DataType;
      static inline void iter4(Neuron &n, DataType d, DataAccumType&) { *d = n.v; }
    };
    
    Layer &fl = *layerFull;
    Layer &bl = *layerPre;

    iterateNeurons2<IL>(fl.layout, fl.layout, fl.neurons, src);
    fl.passFull(&bl, threadsCount);
    iterateNeurons2<IS>(bl.layout, bl.layout, bl.neurons, dst);
  }
  
  
  bool loadImage(int fromIndex, int toIndex) {
    unsigned char *src = data.data();
    if (!layerPre) src += toIndex*imageSize;
    
    fseeko64(f, fromIndex*imageSize, SEEK_SET);
    if (!fread(src, imageSize, 1, f))
      return fclose(f), f = nullptr, false;
    
    if (layerPre) preprocess(src, values.data() + toIndex*preparedImageSize);
    
    return true;
  }
  
  
  bool loadImages() {
    for(int i = 0; i < imagesInMemory; ++i) {
      int j = rand()%imagesInFile;
      if (i != j) std::swap(shuffle[i], shuffle[j]);
    }
    
    typedef std::pair<int, int> Pair;
    typedef std::set<Pair> Set;
    Set set;
    for(int i = 0; i < imagesInMemory; ++i)
      set.insert(Pair(shuffle[i], i));
    for(Set::iterator i = set.begin(); i != set.end(); ++i)
      loadImage(i->first, i->second);
    
    return true;
  }
  
  
  void prepareMeasure() {
    if (measuresPerBlock <= 0) return;
    int sy = segment->sy;
    int sx = segment->sx;
    int sz = segment->sz;
    int sxz = sx*sz;
    int w = (layerPre ? layerPre : layerFull)->layout.getW();
    int h = (layerPre ? layerPre : layerFull)->layout.getH();
    int rowstride = w*sz;
    NeuronReal *dst = valuesMeasure.data();
    for(int i = 0; i < measuresPerBlock; ++i) {
      int index = rand()%imagesInMemory;
      int x = rand()%(w - sx + 1);
      int y = rand()%(h - sy + 1);
      if (layerPre) {
        const NeuronReal *src = values.data() + index*preparedImageSize + y*rowstride + x*sz;
        for(int j = 0; j < sy; ++j, src += rowstride, dst += sxz)
          memcpy(dst, src, sxz*sizeof(*dst));
      } else {
        const unsigned char *src = data.data() + index*preparedImageSize + y*rowstride + x*sz;
        for(int j = 0; j < sy; ++j, src += rowstride - sxz)
        for(int k = 0; k < sxz; ++k, ++src, ++dst)
          *dst = *src/(NeuronReal)255;
      }
    }
  }
  
  
  bool prepare() override {
    assert(infile);
    assert(layerFull);
    assert(loadImagesCount > 0);

    Layer &fl = layerFull->front();
    Layer &bl = layerFull->back();
    
    imageSize = fl.layout.getActiveCount();
    f = fopen(infile, "rb");
    if (!f) return false;
    fseeko64(f, 0, SEEK_END);
    imagesInFile = ftello64(f)/imageSize;
    if (imagesInFile < 1) return fclose(f), f = nullptr, false;
    imagesInMemory = loadImagesCount > imagesInFile ? imagesInFile : loadImagesCount;
    
    for(Layer *l = layerFull; l; l = l->next)
      l->split(threadsCount);
    
    Layout l = layerPre ? layerPre->layout : layerFull->layout;
    assert(l.getW() >= segment->sx);
    assert(l.getH() >= segment->sy);
    assert(l.getD() == segment->sz);
    
    measureLayout = Layout(segment->sx, segment->sy, segment->sz);
    valuesMeasure.resize(measuresPerBlock * measureLayout.getActiveCount());
    if (layerPre) {
      assert(l);
      preparedImageSize = layerPre->layout.getActiveCount();
      trainLayout = Layout(l.getW(), l.getH(), l.getD());
      data.resize(imageSize);
      values.resize(imagesInMemory * preparedImageSize);
    } else {
      trainLayout = measureLayout;
      data.resize(imagesInMemory * imageSize);
      values.resize(segment->sx * segment->sy * segment->sz);
    }

    segment->f_values = values.data();
    segment->layout = trainLayout;
    tmpdata.resize(bl.layout.getActiveCount());
    if (tmpdata.size() < imageSize) tmpdata.resize(imageSize);
    
    size_t memsize = data.size()*sizeof(data.front())
                   + values.size()*sizeof(values.front())
                   + valuesMeasure.size()*sizeof(valuesMeasure.front())
                   + tmpdata.size()*sizeof(tmpdata.front());
    printf("allocated size: %lld\n", (long long)(memsize));

    shuffle.resize(imagesInFile);
    for(int i = 0; i < imagesInFile; ++i)
      shuffle[i] = i;

    if (!loadImages()) return false;
    prepareMeasure();
    return true;
  }
  
  
  void finish() override
    { if (f) fclose(f), f = nullptr; }

    
  bool prepareBlock(int block, bool measureOnly) override {
    if (block > 0 && blocksPerLoading > 0 && (block % blocksPerLoading) == 0 && !loadImages())
      return false;
    seed = rand();
    return true;
  }
  
  
  void finishBlock(int block) override {
    if (outfile) {
      struct IL: public Iter {
        typedef const unsigned char* DataType;
        static inline void iter4(Neuron &n, DataType d, DataAccumType&) { n.v = *d/(NeuronReal)255; }
      };
      struct IS: public Iter {
        typedef unsigned char* DataType;
        static inline void iter4(Neuron &n, DataType d, DataAccumType&) { *d = n.v < 0 ? 0 : n.v > 1 ? 255 : (unsigned char)(n.v*255.999); }
      };
      
      Layer &fl = *layerFull;
      Layer &bl = fl.back();

      std::string outfile0(outfile);
      std::string outfile1 = outfile0 + ".1.tga";
      outfile0 += ".0.tga";

      int index = rand()%imagesInFile;
      fseeko64(f, index*imageSize, SEEK_SET);
      fread(tmpdata.data(), imageSize, 1, f);
      tgaSave(outfile0.c_str(), tmpdata.data(), fl.layout.getW(), fl.layout.getH(), fl.layout.getD());

      iterateNeurons2<IL>(fl.layout, fl.layout, fl.neurons, tmpdata.data());
      fl.passFull(&bl, threadsCount);
      
      iterateNeurons2<IS>(bl.layout, bl.layout, bl.neurons, tmpdata.data());
      tgaSave(outfile1.c_str(), tmpdata.data(), bl.layout.getW(), bl.layout.getH(), bl.layout.getD());
      
      segment->saveDemo();
    }
  }


  void loadData(Barrier &barrier, int block, int iter, bool measureOnly) override {
    int tid = barrier.tid;
    int threads = barrier.threads;
    int sx = segment->sx;
    int sy = segment->sy;
    int sz = segment->sz;
    int sxz = sx*sz;

    if (measureOnly) {
      if (!tid) {
        segment->layout = measureLayout;
        segment->f_values = valuesMeasure.data() + iter*sy*sxz;
        x = y = z = 0;
      }
    } else
    if (layerPre) {
      if (!tid) {
        unsigned int s = randomNext(seed & iter);
        int index = (s = randomNext(s))%imagesInMemory;
        x = (s = randomNext(s)) % (layerPre->layout.getW() - sx + 1);
        y = (s = randomNext(s)) % (layerPre->layout.getH() - sy + 1);
        z = 0;
        segment->layout = trainLayout;
        segment->f_values = values.data() + index*preparedImageSize;
      }
    } else {
      int w = layerFull->layout.getW();
      int h = layerFull->layout.getH();
      
      unsigned int s = randomNext(seed & iter);
      int index = (s = randomNext(s))%imagesInMemory;
      int x0    = (s = randomNext(s))%(w - sx + 1);
      int y0    = (s = randomNext(s))%(h - sy + 1);
      
      int rowstride = w*sz;
      int dr = rowstride*threads - sxz;
      int vdr = sxz*(threads - 1);

      const unsigned char *id0 = data.data() + index*imageSize + y0*rowstride + x0*sz;
      const unsigned char *id = id0 + tid*rowstride;
      NeuronReal *iv = values.data() + tid*sxz;
      
      for(const unsigned char *e = id0 + sy*rowstride; id < e; id += dr, iv += vdr)
      for(const unsigned char *e = id + sxz;        id < e; ++id, ++iv)
        *iv = *id/(NeuronReal)255;
      
      if (!tid) {
        segment->layout = trainLayout;
        segment->f_values = values.data();
        x = 0, y = 0, z = 0;
      }
    }
  }
};


#endif