#ifndef NNLAYER3_MT_INC_CPP
#define NNLAYER3_MT_INC_CPP


#include "nnlayer3.inc.cpp"


#include <atomic>
#include <thread>
#include <vector>


class Barrier {
private:
  std::atomic<unsigned int> &counter;
  const unsigned int threads;
  unsigned int next;
public:
  inline Barrier(std::atomic<unsigned int> &counter, unsigned int threads): counter(counter), threads(threads), next() { }
  inline void wait() { next += threads; ++counter; while(counter < next); }
  inline void subwait(int tid) { while(counter < next + tid); }
};


class TrainMT {
private:
  struct LDesc {
    int y0, y1;
    double sumQ;
    LDesc(): y0(), y1(), sumQ() { }
  };

public:
  Layer *layer;
  const unsigned char *dataX;
  const unsigned char *dataY;
  int strideX;
  int strideY;
  int *shuffle;
  int count;
  Real trainRatio;

  TrainMT():
    layer(),
    dataX(),
    dataY(),
    strideX(),
    strideY(),
    shuffle(),
    count(),
    trainRatio() { }

private:
  void trainFunc(int tid, int threads, std::atomic<unsigned int> &barrierCounter, LDesc *ldescs) {
    Barrier barrier(barrierCounter, threads);

    Layer &fl = *layer;
    Layer &bl = layer->back();
    int layersCount = fl.totalLayers();
    LDesc *fld = ldescs, *bld = fld + layersCount - 1;

    Real trainRatio = this->trainRatio;

    int fsxz = fl.sx*fl.sz;
    int bsxz = bl.sx*bl.sz;

    //barrier.subwait(tid);
    //for(LDesc *ld = fld; ld <= bld; ++ld)
    //  printf("t%d %d %d %d\n", tid, (int)(ld-fld), ld->y0, ld->y1);
    //barrier.wait();

    const unsigned char *dataX = this->dataX + fsxz*fld->y0;
    const unsigned char *dataY = this->dataY + bsxz*bld->y0;

    double sumQ = 0;
    for(int i = 0; i < count; ++i) {
      int ii = shuffle[i];
      const unsigned char *curX = dataX + strideX*ii;
      const unsigned char *curY = dataY + strideY*ii;

      barrier.wait();
      const unsigned char *px = curX;
      for(Neuron *in = fl.neurons + fsxz*fld->y0, *e = fl.neurons + fsxz*fld->y1; in < e; ++in, ++px)
        in->v = Real(*px)*Real(1/255.0);

      LDesc *ld = fld + 1;
      for(Layer *l = fl.next; l; l = l->next, ++ld) {
        barrier.wait();
        l->pass(ld->y0, ld->y1);
      }

      double q = 0;
      const unsigned char *py = curY;
      for(Neuron *in = bl.neurons + bsxz*bld->y0, *e = bl.neurons + bsxz*bld->y1; in < e; ++in, ++py) {
        Real v = (in->v - 0.25)*2;
        Real d = Real(*py)*Real(1/255.0) - v;
        in->d *= d * trainRatio;
        d *= d;
        q += d*d;
      }
      sumQ += q;

      if (trainRatio > 0) {
        ld = bld;
        for(Layer *l = &bl; l->prev; l = l->prev, --ld) {
          if (!l->prev->prev) {
            barrier.wait();
            l->backpassWeights(ld->y0, ld->y1);
            break;
          } else
          if (l->next) {
            barrier.wait();
            l->backpassTpl<true>(ld->y0, ld->y1);
            //l->backpassTpl<false>(ld->y0, ld->y1);
            //barrier.wait();
            //l->next->backpassWeights(ld[1].y0, ld[1].y1);
          }
        }
      }

      //if (!tid) printf(" - %d, %f, %f\n", i, q, sumQ);
    }

    ldescs->sumQ = sumQ;
  }

public:
  double train(int threads) {
    assert(threads > 0);
    assert(layer && !layer->prev);
    assert(dataX && dataY && shuffle);
    assert(count > 0);
    assert(trainRatio >= 0);

    int layersCount = layer->totalLayers();
    assert(layersCount > 0);
    std::vector<LDesc> ldescs( threads*layersCount );

    int layerId = 0;
    for(Layer *l = layer; l; l = l->next, ++layerId) {
      assert(layerId < layersCount);
      int tsy = l->sy/threads;
      for(int tid = 0; tid < threads; ++tid) {
        LDesc &desc = ldescs[tid*layersCount + layerId];
        desc.y0 = tid*tsy;
        desc.y1 = desc.y0 + tsy;
        if (tid == threads-1) desc.y1 = l->sy;
      }
    }
    assert(layerId == layersCount);

    std::atomic<unsigned int> barrierCounter(0);
    std::vector<std::thread*> t(threads - 1);
    for(int i = 1; i < threads; ++i)
      t[i-1] = new std::thread(&TrainMT::trainFunc, this, i, threads, std::ref(barrierCounter), &ldescs[i*layersCount]);
    trainFunc(0, threads, barrierCounter, &ldescs[0]);

    double result = ldescs[0].sumQ;
    for(int i = 1; i < threads; ++i)
      { t[i-1]->join(); delete t[i-1]; result += ldescs[i*layersCount].sumQ; }

    return sqrt(sqrt( result/(count * layer->back().countNeurons()) ));
  }
};


#endif