#ifndef LAYER_CONV_SHARED_INC_CPP
#define LAYER_CONV_SHARED_INC_CPP



#include "layer.conv.inc.cpp"



template<typename Iter>
void iterateTestConvolutionShared(Layout cl, Layout pl, Kernel k, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  if (!cl) return;
  assert(pl);
  assert(k);
  assert(c_neurons);
  assert(p_neurons);
  assert(weights);
  assert(pl.x0 + k.ox >= 0 && pl.x0 + (cl.getW()-1)*k.dx + k.ox + k.sx <= pl.sx);
  assert(pl.y0 + k.oy >= 0 && pl.y0 + (cl.getH()-1)*k.dy + k.oy + k.sy <= pl.sy);
 
  for(int cy = cl.y0; cy < cl.y1; ++cy)
  for(int cx = cl.x0; cx < cl.x1; ++cx)
  for(int cz = cl.z0; cz < cl.z1; ++cz) {
    int ci = (cy*cl.sx + cx)*cl.sz + cz;
    Neuron &cn = c_neurons[ci];
    typename Iter::AccumType a = {};
    Iter::init(cn, a);
    
    for(int ky = 0; ky < k.sy; ++ky)
    for(int kx = 0; kx < k.sx; ++kx)
    for(int pz = pl.z0; pz < pl.z1; ++pz) {
      int wi = (ky*k.sx + kx)*pl.getD() + pz - pl.z0;
      Weight &w = weights[wi];

      int px = pl.x0 + (cx - cl.x0)*k.dx + k.ox + kx;
      int py = pl.y0 + (cy - cl.y0)*k.dy + k.oy + ky;
      int pi = (py*pl.sx + px)*pl.sz + pz;
      Neuron &pn = p_neurons[pi];
      
      Iter::iter(pn, w, a);
    }
    
    Iter::done(cn, a);
  }
}


template<typename Iter>
void iterateConvolutionSharedDyn(Layout cl, Layout pl, Layout wl, Kernel k, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  if (!cl) return;
  assert(pl);
  assert(wl);
  assert(k);
  assert(c_neurons);
  assert(p_neurons);
  assert(weights);
  assert(cl.isSubLayoutOf(wl));
  assert(pl.x0 + k.ox >= 0 && pl.x0 + (wl.getW()-1)*k.dx + k.ox + k.sx <= pl.sx);
  assert(pl.y0 + k.oy >= 0 && pl.y0 + (wl.getH()-1)*k.dy + k.oy + k.sy <= pl.sy);

  int c_h    = cl.getH();
  int c_w    = cl.getW();
  int c_d    = cl.getD();
  int c_swz  = c_w*cl.sz;
  int c_shxz = c_h*cl.sx*cl.sz;
  int c_dx   = cl.sz - c_d;
  int c_dy   = (cl.sx - c_w)*cl.sz;

  int p_d    = pl.getD();
  int p_dx   = k.dx*pl.sz;
  int p_dy   = k.dy*pl.sx*pl.sz - c_w*p_dx;

  int k_sxd  = k.sx*p_d;
  int p_ddy  = (pl.sx - k.sx)*pl.sz;
  int p_ddx  = pl.sz - p_d;

  Neuron *icn = c_neurons + (cl.y0*cl.sx + cl.x0)*cl.sz + cl.z0;
  Neuron *ipn = p_neurons + ((pl.y0 + (cl.y0 - wl.y0)*k.dy + k.oy)*pl.sx + pl.x0 + (cl.x0 - wl.x0)*k.dx + k.ox)*pl.sz + pl.z0;
  Weight *ew = weights + k.sy*k_sxd;

  for(Neuron *e = icn + c_shxz; icn < e; icn += c_dy, ipn += p_dy)
  for(Neuron *e = icn +  c_swz; icn < e; icn += c_dx, ipn += p_dx)
  for(Neuron *e = icn +    c_d; icn < e; ++icn) {
    typename Iter::AccumType a;
    Iter::init(*icn, a);

    Neuron *iipn = ipn;
    for(Weight *iw =    weights; iw < ew; iipn += p_ddy)
    for(Weight *e = iw +  k_sxd; iw < e;  iipn += p_ddx)
    for(Weight *e = iw +    p_d; iw < e;  ++iw, ++iipn)
      Iter::iter(*iipn, *iw, a);

    Iter::done(*icn, a);
  }
}


template<typename Iter, int KSX, int KSY, int PD>
void iterateConvolutionSharedXYD(Layout cl, Layout pl, Layout wl, Kernel k, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  if (!cl) return;
  assert(pl);
  assert(wl);
  assert(k);
  assert(c_neurons);
  assert(p_neurons);
  assert(weights);
  assert(cl.isSubLayoutOf(wl));
  assert(pl.x0 + k.ox >= 0 && pl.x0 + (wl.getW()-1)*k.dx + k.ox + k.sx <= pl.sx);
  assert(pl.y0 + k.oy >= 0 && pl.y0 + (wl.getH()-1)*k.dy + k.oy + k.sy <= pl.sy);
  
  assert(KSX == k.sx);
  assert(KSY == k.sy);
  assert(PD == pl.getD());

  int c_h    = cl.getH();
  int c_w    = cl.getW();
  int c_d    = cl.getD();
  int c_swz  = c_w*cl.sz;
  int c_shxz = c_h*cl.sx*cl.sz;
  int c_dx   = cl.sz - c_d;
  int c_dy   = (cl.sx - c_w)*cl.sz;

  int p_dx   = k.dx*pl.sz;
  int p_dy   = k.dy*pl.sx*pl.sz - c_w*p_dx;

  int p_ddy  = (pl.sx - KSX)*pl.sz;
  int p_ddx  = pl.sz - PD;

  Neuron *icn = c_neurons + (cl.y0*cl.sx + cl.x0)*cl.sz + cl.z0;
  Neuron *ipn = p_neurons + ((pl.y0 + (cl.y0 - wl.y0)*k.dy + k.oy)*pl.sx + pl.x0 + (cl.x0 - wl.x0)*k.dx + k.ox)*pl.sz + pl.z0;

  for(Neuron *e = icn + c_shxz; icn < e; icn += c_dy, ipn += p_dy)
  for(Neuron *e = icn +  c_swz; icn < e; icn += c_dx, ipn += p_dx)
  for(Neuron *e = icn +    c_d; icn < e; ++icn) {
    typename Iter::AccumType a;
    Iter::init(*icn, a);

    Neuron *iipn = ipn;
    Weight *iw = weights;
    for(int i = 0; i < KSY; ++i, iipn += p_ddy)
    for(int i = 0; i < KSX; ++i, iipn += p_ddx)
    for(int i = 0; i <  PD; ++i, ++iw, ++iipn)
      Iter::iter(*iipn, *iw, a);

    Iter::done(*icn, a);
  }
}

typedef void (*iterateConvolutionSharedFunc)(Layout, Layout, Layout, Kernel, Neuron*, Neuron*, Weight*);
template<typename Iter, int KSX, int KSY>
iterateConvolutionSharedFunc getIterateConvolutionSharedFuncXY(int pd) {
  if (pd <= 8) switch(pd) {
  case 1:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 1>;
  case 2:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 2>;
  case 3:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 3>;
  case 4:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 4>;
  case 5:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 5>;
  case 6:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 6>;
  case 7:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 7>;
  case 8:  return &iterateConvolutionSharedXYD<Iter, KSX, KSY, 8>;
  }
  return &iterateConvolutionSharedDyn<Iter>;
}


template<typename Iter>
iterateConvolutionSharedFunc getIterateConvolutionSharedFunc(int ksx, int ksy, int pd) {
  if (0 && ksx == ksy && pd <= 8) switch(ksx) {
  case 1:  return getIterateConvolutionSharedFuncXY<Iter, 1, 1>(pd);
  case 2:  return getIterateConvolutionSharedFuncXY<Iter, 2, 2>(pd);
  case 3:  return getIterateConvolutionSharedFuncXY<Iter, 3, 3>(pd);
  case 4:  return getIterateConvolutionSharedFuncXY<Iter, 4, 4>(pd);
  case 5:  return getIterateConvolutionSharedFuncXY<Iter, 5, 5>(pd);
  case 6:  return getIterateConvolutionSharedFuncXY<Iter, 6, 6>(pd);
  case 7:  return getIterateConvolutionSharedFuncXY<Iter, 7, 7>(pd);
  case 8:  return getIterateConvolutionSharedFuncXY<Iter, 8, 8>(pd);
  }
  return &iterateConvolutionSharedDyn<Iter>;
}


template<typename Iter>
void iterateConvolutionShared(const Layout &cl, const Layout &pl, const Layout &wl, const Kernel &k, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  iterateConvolutionSharedFunc f = getIterateConvolutionSharedFunc<Iter>(k.sx, k.sy, pl.getD());
  f(cl, pl, wl, k, c_neurons, p_neurons, weights);
}




template<typename Iter>
void iterateConvolutionSharedPoint(Layout cl, Layout pl, Layout wl, Kernel k, int kx, int ky, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  if (!cl) return;
  assert(pl);
  assert(wl);
  assert(k);
  assert(c_neurons);
  assert(p_neurons);
  assert(weights);
  assert(cl.isSubLayoutOf(wl));
  assert(kx >= 0 && kx < k.sx);
  assert(ky >= 0 && ky < k.sy);
  assert(pl.x0 + k.ox >= 0 && pl.x0 + (wl.getW()-1)*k.dx + k.ox + k.sx <= pl.sx);
  assert(pl.y0 + k.oy >= 0 && pl.y0 + (wl.getH()-1)*k.dy + k.oy + k.sy <= pl.sy);

  int c_h    = cl.getH();
  int c_w    = cl.getW();
  int c_d    = cl.getD();
  int c_swz  = c_w*cl.sz;
  int c_shxz = c_h*cl.sx*cl.sz;
  int c_dx   = cl.sz - c_d;
  int c_dy   = (cl.sx - c_w)*cl.sz;

  int p_d    = pl.getD();
  int p_dx   = k.dx*pl.sz;
  int p_dy   = k.dy*pl.sx*pl.sz - c_w*p_dx;

  Neuron *icn = c_neurons + (cl.y0*cl.sx + cl.x0)*cl.sz + cl.z0;
  Neuron *ipn = p_neurons + ((pl.y0 + (cl.y0 - wl.y0)*k.dy + k.oy + ky)*pl.sx + pl.x0 + (cl.x0 - wl.x0)*k.dx + k.ox + kx)*pl.sz + pl.z0;
  weights += (ky*k.sx + kx)*p_d;
  Weight *ew = weights + p_d;

  for(Neuron *e = icn + c_shxz; icn < e; icn += c_dy, ipn += p_dy)
  for(Neuron *e = icn +  c_swz; icn < e; icn += c_dx, ipn += p_dx)
  for(Neuron *e = icn +    c_d; icn < e; ++icn,       ipn -= p_d)
  for(Weight *iw = weights; iw < ew; ++ipn, ++iw)
    Iter::iter2(*icn, *ipn, *iw);
}




class LayerConvSharedBase: public Layer {
public:
  std::vector<Weight> mtWeights;
  
  using Layer::Layer;
  

  void split(int threadsCount) override {
    Layer::split(threadsCount);
    Weight w = {};
    mtWeights.clear();
    mtWeights.resize(threadsCount*weightsCount, w);
  }

  
  inline void sumWeights(int tid, int threads) {
    int wc = weightsCount;
    Weight *iw = weights + tid;
    Weight *ia = mtWeights.data() + tid;
    Weight *ea = mtWeights.data() + threads*wc;
    for(Weight *ew = weights + wc; iw < ew; iw += threads, ia += threads) {
      WeightReal w = iw->w;
      for(Weight *iia = ia; iia < ea; iia += wc)
        w += iia->w, iia->w = 0;
      iw->w = w;
    }
  }
};



template<Func func>
class LayerConvShared: public LayerConvSharedBase {
public:
  Kernel kernel;


  LayerConvShared(Layer &prev, const Layout &layout, const Kernel &kernel, Weight *weights = nullptr):
    LayerConvSharedBase(&prev, layout, kernel.sx*kernel.sy*prev.back().layout.getD(), weights),
    kernel(kernel)
  {
    assert(kernel);
    stat.links = weightsCount*neuronsCount;
    if (ownWeights) fillWeights(-1, 1);
  }

  
  void pass(Barrier &barrier) override {
    struct I: public Iter {
      static inline void init(Neuron&, AccumType &a) { a.v = 0; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { a.v += n.v * w.w; }
      static inline void done(Neuron &n, AccumType &a) { func(n, a.v); }
    };
    iterateConvolutionShared<I>(mtLayouts[barrier.tid], prev->layout, layout, kernel, neurons, prev->neurons, weights);
  }
  

  void backpassWeights(Barrier &barrier) override {
    struct I: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = n.d; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { w.w += n.v * a.v; }
    };
    iterateConvolutionShared<I>(mtLayouts[barrier.tid], prev->layout, layout, kernel, neurons, prev->neurons, &mtWeights[barrier.tid * weightsCount]);
    barrier.wait();
    sumWeights(barrier.tid, barrier.threads);
  }
  
  
  void backpassDeltas(Barrier &barrier) override {
    struct I: public Iter {
      static inline void iter2(Neuron &cn, Neuron &pn, Weight &w) { pn.a.v += cn.d * w.w; }
      static inline void iter3(Neuron &n) { n.d *= n.a.v; n.a.v = 0; }
    };
    int ksx = kernel.sx, ksy = kernel.sy;
    for(int kx = 0; kx < ksx; ++kx)
    for(int ky = 0; ky < ksy; ++ky) {
      iterateConvolutionSharedPoint<I>(mtLayouts[barrier.tid], prev->layout, layout, kernel, kx, ky, neurons, prev->neurons, weights);
      barrier.wait();
    }
    iterateNeurons<I>(mtPrevLayouts[barrier.tid], prev->neurons);
  }
  
  
  void testPass() override {
    struct I: public Iter {
      static inline void init(Neuron&, AccumType &a) { a.v = 0; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { a.v += n.v * w.w; }
      static inline void done(Neuron &n, AccumType &a) { func(n, a.v); }
    };
    iterateTestConvolutionShared<I>(layout, prev->layout, kernel, neurons, prev->neurons, weights);
  }

    
  void testBackpass() override {
    struct I: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = n.d; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { n.a.v += a.v * w.w; }
      static inline void iter3(Neuron &n) { n.d *= n.a.v; n.a.v = 0; }
    };
    struct IW: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = n.d; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { w.w += a.v * n.v; }
    };
    clearAccum();
    iterateTestConvolutionShared<I>(layout, prev->layout, kernel, neurons, prev->neurons, weights);
    iterateTestConvolutionShared<IW>(layout, prev->layout, kernel, neurons, prev->neurons, weights);
    iterateNeurons<I>(prev->layout, prev->neurons);
    clearAccum();
  }
};



template<Func func>
class LayerDeconvShared: public LayerConvSharedBase {
public:
  Kernel kernel;


  LayerDeconvShared(Layer &prev, const Layout &layout, const Kernel &kernel, Weight *weights = nullptr):
    LayerConvSharedBase(&prev, layout, kernel.sx*kernel.sy*layout.getD(), weights),
    kernel(kernel)
  {
    assert(kernel);
    stat.links = weightsCount*neuronsCount;
    if (ownWeights) fillWeights(-1, 1);
  }

  
  void pass(Barrier &barrier) override {
    struct I: public Iter {
      static inline void iter2(Neuron &cn, Neuron &pn, Weight &w) { pn.a.v += cn.v * w.w; }
      static inline void iter3(Neuron &n) { func(n, n.a.v); n.a.v = 0; }
    };
    int k_sx = kernel.sx, k_sy = kernel.sy;
    for(int kx = 0; kx < k_sx; ++kx)
    for(int ky = 0; ky < k_sy; ++ky) {
      iterateConvolutionSharedPoint<I>(mtPrevLayouts[barrier.tid], layout, prev->layout, kernel, kx, ky, prev->neurons, neurons, weights);
      barrier.wait();
    }
    iterateNeurons<I>(mtLayouts[barrier.tid], neurons);
  }
  
  
  void backpassWeights(Barrier &barrier) override {
    struct I: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = n.v; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { w.w += n.d * a.v; }
    };
    iterateConvolutionShared<I>(mtPrevLayouts[barrier.tid], layout, prev->layout, kernel, prev->neurons, neurons, &mtWeights[barrier.tid * weightsCount]);
    barrier.wait();
    sumWeights(barrier.tid, barrier.threads);
  }
  
  
  void backpassDeltas(Barrier &barrier) override {
    struct I: public Iter {
      static inline void init(Neuron&, AccumType &a) { a.v = 0; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { a.v += n.d * w.w; }
      static inline void done(Neuron &n, AccumType &a) { n.d *= a.v; }
    };
    iterateConvolutionShared<I>(mtPrevLayouts[barrier.tid], layout, prev->layout, kernel, prev->neurons, neurons, weights);
  }

  
  void testPass() override {
    struct I: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = n.v; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { n.a.v += a.v * w.w; }
      static inline void iter3(Neuron &n) { func(n, n.a.v); n.a.v = 0; }
    };
    clearAccum();
    iterateTestConvolutionShared<I>(prev->layout, layout, kernel, prev->neurons, neurons, weights);
    iterateNeurons<I>(layout, neurons);
    clearAccum();
  }
  
  
  void testBackpass() override {
    struct I: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = 0; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { a.v += n.d * w.w; }
      static inline void done(Neuron &n, AccumType &a) { n.d *= a.v; }
    };
    struct IW: public Iter {
      static inline void init(Neuron &n, AccumType &a) { a.v = n.v; }
      static inline void iter(Neuron &n, Weight &w, AccumType &a) { w.w += n.d * a.v; }
    };
    iterateTestConvolutionShared<I>(prev->layout, layout, kernel, prev->neurons, neurons, weights);
    iterateTestConvolutionShared<IW>(prev->layout, layout, kernel, prev->neurons, neurons, weights);
  }
};

#endif