#ifndef TRAIN_IMAGE_INC_CPP
#define TRAIN_IMAGE_INC_CPP


#include <set>

#include "train.inc.cpp"
#include "layer.simple.inc.cpp"


class TrainerImage: public Trainer {
protected:
  std::vector<unsigned char> data;
  std::vector<unsigned char> tmpdata;
  std::vector<int> shuffle;
  std::vector<int> shuffle2;
  Layout pbl;
  Layout::List flist, blist;
  FILE *f;
  size_t imgsize;
  int count;
  int workCount;

public:
  int pad;
  const char *datafile;
  const char *outfile;
  Layer *dataLayer;
  
  TrainerImage(): f(), imgsize(), count(), workCount(), pad(), datafile(), outfile(), dataLayer() { }

protected:
  bool prepare() override {
    assert(datafile);   
    assert(fl->layout.getD() == 3);
    
    #ifndef NDEBUG
    Layer *dl = dataLayer ? dataLayer : fl;
    assert(dl->layout.getW() == bl->layout.getW());
    assert(dl->layout.getH() == bl->layout.getH());
    assert(dl->layout.getD() == bl->layout.getD());
    #endif
    
    imgsize = fl->layout.getActiveCount();
    fl->layout.split(flist, threadsCount);
    bl->layout.split(blist, threadsCount);
    pbl = bl->layout;
    pbl.padXY(pad);
    
    f = fopen(datafile, "rb");
    if (!f) return false;
    
    fseeko64(f, 0, SEEK_END);
    long long size = ftello64(f);
    count = size/imgsize;
    if (count < 1) return fclose(f), f = nullptr, false;

    workCount = itersPerBlock > count ? count : itersPerBlock;
    printf("allocated size: %lld\n", (long long)(imgsize*workCount));
    data.resize(workCount*imgsize);

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

    shuffle2.resize(workCount);
    for(int i = 0; i < workCount; ++i)
      shuffle2[i] = i;
    
    return loadBlocks();
    //return true;
  }
  
  
  void finish() override
    { if (f) fclose(f), f = nullptr; }

    
  bool loadBlocks() {
    for(int i = 0; i < workCount; ++i) {
      int j = rand()%count;
      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 < workCount; ++i)
      set.insert(Pair(shuffle[i], i));
    
    for(Set::iterator i = set.begin(); i != set.end(); ++i) {
      fseeko64(f, i->first*imgsize, SEEK_SET);
      if (!fread(data.data() + i->second*imgsize, imgsize, 1, f))
        return fclose(f), f = nullptr, false;
    }
    
    return true;
  }

  bool prepareBlock() override {
    for(int i = 0; i < workCount; ++i) {
      int j = rand()%workCount;
      if (i != j) std::swap(shuffle2[i], shuffle2[j]);
    }
    //return loadBlocks();
    return true;
  }
  
  
  void finishBlock() override {
    if (outfile && !dataLayer) {
      std::string outfile0(outfile);
      std::string outfile1 = outfile0 + ".1.tga";
      outfile0 += ".0.tga";
      
      unsigned char *id0 = data.data() + shuffle2[(itersPerBlock-1)%workCount]*imgsize;
      tgaSave(outfile0.c_str(), id0, fl->layout.getW(), fl->layout.getH(), fl->layout.getD());

      struct I: 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); }
      };
      
      tmpdata.resize(imgsize);
      unsigned char *id1 = tmpdata.data();
      iterateNeurons2<I>(bl->layout, bl->layout, bl->neurons, id1);
      tgaSave(outfile1.c_str(), id1, bl->layout.getW(), bl->layout.getH(), bl->layout.getD());
    }
  }


  void loadData(Barrier &barrier, int, int iter) override {
    struct I: public Iter {
      typedef const unsigned char* DataType;
      static inline void iter4(Neuron &n, DataType d, DataAccumType&) { n.v = *d/(NeuronReal)255; }
    };
    const unsigned char *id = data.data() + shuffle2[iter%workCount]*imgsize;
    iterateNeurons2<I>(flist[barrier.tid], fl->layout, fl->neurons, id);
  }


  Quality verifyData(Barrier &barrier, int, int iter) override {
    Layout l = blist[barrier.tid];
    Layout dl = bl->layout;
    Layout pl = pbl;
    
    int d = l.getD();
    int w = l.getW();
    int dx = l.sz - d;
    int dy = (l.sx - w)*l.sz;
    int ddx = dl.getD();
    int ddy = (dl.getW() - w)*ddx;
      
    AccumReal aq = 0;
    NeuronReal ratio = this->ratio;
    Neuron *in = bl->neurons + (l.y0*l.sx + l.x0)*l.sz + l.z0;
    const unsigned char *id = data.data() + shuffle2[iter%workCount]*imgsize + ((l.y0-dl.y0)*l.sx + l.x0-dl.x0)*l.sz + l.z0-dl.z0;
    
    for(int y = l.y0; y < l.y1; ++y, in += dy, id += ddy) {
      bool outside = y < pl.y0 || y >= pl.y1;
      for(int x = l.x0; x < l.x1; ++x, in += dx, id += ddx) {
        if (outside || x < pl.x0 || x >= pl.x1) {
          for(Neuron *e = in + d; in < e; ++in) in->d = 0;
        } else {
          const unsigned char *iid = id;
          for(Neuron *e = in + d; in < e; ++in, ++iid) {
            NeuronReal v1 = *iid/(NeuronReal)255;
            NeuronReal v0 = in->v;
            NeuronReal diff = v1 - v0;
            in->d *= diff*ratio;
            aq += diff*diff;
          }
        }
      }
    }
    
    return Quality( sqrt(aq/pbl.getActiveCount()) );
  }
};


#endif