#ifndef NNLAYER2_MT_INC_CPP

#define NNLAYER2_MT_INC_CPP

#include "nnlayer2.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); }

};

class TrainMT {

private:

  struct LDesc {

    int nb, ne, lb, le;

    double sumQ;

    LDesc(): nb(), ne(), lb(), le(), 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;

    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;

      const unsigned char *px = curX;

      for(Neuron *in = fl.neurons + fld->nb, *e = fl.neurons + fld->ne; 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->nb, ld->ne);

      }

      double q = 0;

      const unsigned char *py = curY;

      for(Neuron *in = bl.neurons + bld->nb, *e = bl.neurons + bld->ne; in < e; ++in, ++py) {

        Real d = Real(*py)*Real(1/255.0) - in->v;

        in->d *= d * trainRatio;

        q += d*d;

      }

      sumQ += q;

      if (trainRatio > 0) {

        ld = bld - 1;

        for(Layer *l = bl.prev; l; l = l->prev, --ld) {

          barrier.wait();

          l->backpass(ld->lb, ld->le);

        }

      }

      //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 tsize = l->size/threads;

      for(int tid = 0; tid < threads; ++tid) {

        LDesc &desc = ldescs[tid*layersCount + layerId];

        desc.nb = tid*tsize;

        desc.ne = desc.nb + tsize;

        if (tid == threads-1) desc.ne = l->size;

      }

      if (int lsize = l->size*l->lsize) {

        int tlsize = lsize/threads;

        int ipn = l->links[ l->lfirst ].nprev;

        int tid = 0;

        int count = 0;

        ldescs[tid*layersCount + layerId].lb = l->lfirst;

        if (threads > 1) {

          for(int il = l->lfirst; il != lsize; il = l->links[il].lnext, ++count) {

            Link &link = l->links[il];

            if (ipn != link.nprev) {

              if (count >= tlsize) {

                ldescs[tid*layersCount + layerId].le = il;

                ++tid;

                count -= tlsize;

                ldescs[tid*layersCount + layerId].lb = il;

                if (tid == threads - 1) break;

              }

              ipn = link.nprev;

            }

          }

        }

        ldescs[tid*layersCount + layerId].le = lsize;

      }

    }

    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 result/(count * layer->back().size);

  }

};

#endif