#ifndef TRAIN_SEGMENT_INC_CPP
#define TRAIN_SEGMENT_INC_CPP


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


class TrainerSegment {
private:
  std::atomic<unsigned int> barrierCounter;
  std::vector<QualityPair> qualities;

public:
  Segment *segment;
  AccumReal ratio;
  int threadsCount;
  
  int measuresPerBlock;
  int trainsPerBlock;
  int blocksPerSaving;
  
  int blocksCount;
  AccumReal qmin;

protected:
  volatile int x, y, z;
  
  virtual bool prepare() { return true; }
  virtual bool prepareBlock(int block, bool measureOnly) { return true; }
  virtual void finishBlock(int block) { }
  virtual void finish() { }

  virtual void loadData(Barrier &barrier, int block, int iter, bool measure) { }

private:
  void threadFunc(int tid, unsigned int seed, int block, bool measureOnly) {
    Barrier barrier(barrierCounter, tid, threadsCount, seed);

    QualityPair q;
    if (!measureOnly) {
      for(int i = 0; i < trainsPerBlock; ++i) {
        barrier.wait();
        loadData(barrier, block, i, false);
        barrier.wait();
        q.train += segment->pass(barrier, x, y, z, ratio);
      }
    }
    
    for(int i = 0; i < measuresPerBlock; ++i) {
      barrier.wait();
      loadData(barrier, block, i, true);
      barrier.wait();
      q.measure += segment->pass(barrier, x, y, z, 0);
    }
    
    qualities[tid] = q;
  }


  QualityPair runThreads(int block, bool measureOnly) {
    barrierCounter = 0;
    std::vector<std::thread*> t(threadsCount, nullptr);
    for(int i = 1; i < threadsCount; ++i)
      t[i] = new std::thread(&TrainerSegment::threadFunc, this, i, rand(), block, measureOnly);
    threadFunc(0, rand(), block, measureOnly);

    QualityPair q = qualities[0];
    for(int i = 1; i < threadsCount; ++i)
      { t[i]->join(); delete t[i]; q += qualities[i]; }
    
    q.measure *= 1/(AccumReal)measuresPerBlock;
    q.train *= 1/(AccumReal)trainsPerBlock;
    return q;
  }


public:
  TrainerSegment():
    barrierCounter(0),
    segment(),
    ratio(),
    threadsCount(1),
    measuresPerBlock(),
    trainsPerBlock(),
    blocksPerSaving(),
    blocksCount(),
    qmin(),
    x(), y(), z()
    { }


  QualityPair run() {
    int trainsPerBlock = ratio > 0 ? this->trainsPerBlock : 0;
    int blocksCount = trainsPerBlock > 0 || this->blocksCount > 0 ? this->blocksCount : 1;

    assert(segment);
    assert(threadsCount > 0);
    assert(measuresPerBlock >= 0);
    assert(trainsPerBlock >= 0);
    assert(measuresPerBlock + trainsPerBlock > 0);

    QualityPair bad(Quality::bad(), Quality::bad());
    
    printf("training segment: threads %d, trainsPerBlock %d, measuresPerBlock %d, ratio: %lf\n", threadsCount, trainsPerBlock, measuresPerBlock, ratio);
    fflush(stdout);

    qualities.clear();
    qualities.resize(threadsCount);
    segment->split(threadsCount);
    
    if (!prepare())
      return printf("cannot prepare\n"), bad;

    
    QualityPair result = bad, best = result, saved = result;
    long long fullTimeStartUs = timeUs();
    int i = 0;
    int bps = blocksPerSaving > 0 ? blocksPerSaving : 1;
    int nextSave = i + bps;
    while(true) {
      bool measureOnly = measuresPerBlock > 0 && (!i || trainsPerBlock <= 0);
      
      if (!prepareBlock(i, measureOnly)) {
        printf("cannot prepare block\n");
        result = bad;
        break;
      };

      long long runTimeUs = timeUs();
      result = runThreads(i, measureOnly);
      runTimeUs = timeUs() - runTimeUs;

      finishBlock(i);

      long long t = timeUs();
      long long fullTimeUs = t - fullTimeStartUs;
      fullTimeStartUs = t;
      ++i;

      Quality q = measuresPerBlock > 0 ? result.measure : result.train;
      
      if (i == 1) saved = result;
      bool good = result < best;
      bool done = (blocksCount > 0 && i >= blocksCount) || q.human <= qmin;
      bool saving = !measureOnly && ratio > 0 && (i >= nextSave || done) && result < saved;
      if (good) best = result;
      
      Quality bq = measuresPerBlock > 0 ? best.measure : best.train;
      
      printf("%4d, total %7d, avg.result %12g (%12g), best %12g (%12g), time: %f / %f%s\n",
        i, i*trainsPerBlock,
        q.human, q.train, bq.human, bq.train,
        runTimeUs*0.000001, fullTimeUs*0.000001,
        (saving ? ", saving" : "" ) );
      fflush(stdout);

      if (saving) {
        if (!segment->save()) {
          printf("saving failed\n");
          result = bad;
          break;
        }
        saved = result;
        nextSave += bps;
      }

      if (done) break;
    }

    finish();

    return result;
  }
};


#endif