#ifndef LAYER_SIMPLE_INC_CPP
#define LAYER_SIMPLE_INC_CPP


#include "layer.inc.cpp"


typedef void Func(Neuron &n, AccumReal s);


inline void funcSigmoidExp(Neuron &n, AccumReal s) {
  //if (s > 5) s = 5; else if (s < -5) s = -5;
  AccumReal ss = 1/(1 + std::exp(-s)); n.v = ss; n.d = ss * (1-ss);
}


template<typename Iter>
inline void iterateNeurons(const Layout &l, Neuron *neurons) {
  if (!l) return;
  assert(neurons);
  
  int h = l.y1 - l.y0;
  int w = l.x1 - l.x0;
  int d = l.z1 - l.z0;
  int sz = l.sz;
  int sxz = l.sx*sz;
  int swz = w*sz;
  int shxz = h*sxz;
  int dy = sxz - swz;
  int dx = sz - d;

  Neuron *in = neurons + l.y0*sxz + l.x0*sz + l.z0;

  for(Neuron *e = in + shxz; in < e; in += dy)
  for(Neuron *e = in +  swz; in < e; in += dx)
  for(Neuron *e = in +    d; in < e; ++in)
    Iter::iter3(*in);
}


template<typename Iter>
inline void iterateNeurons2(Layout l, Layout dl, Neuron *neurons, typename Iter::DataType data, int stride = 1, typename Iter::DataAccumType *accum = nullptr) {
  if (!l) return;
  assert(dl);
  assert(neurons);
  assert(l.isSubLayoutOf(dl));
  
  int h    = l.getH();
  int w    = l.getW();
  int d    = l.getD();
  int sxz  = l.sx*l.sz;
  int swz  = w*l.sz;
  int shxz = h*sxz;
  int dy   = sxz - swz;
  int dx   = l.sz - d;

  int d_w  = dl.getW();
  int d_d  = dl.getD();
  int d_dx = (d_d - d)*stride;
  int d_dy = (d_w - w)*d_d*stride;

  Neuron *in = neurons + l.y0*sxz + l.x0*l.sz + l.z0;
  data += (((l.y0 - dl.y0)*d_w + l.x0 - dl.x0)*d_d + l.z0 - dl.z0)*stride;

  for(Neuron *e = in + shxz; in < e; in += dy, data += d_dy)
  for(Neuron *e = in +  swz; in < e; in += dx, data += d_dx)
  for(Neuron *e = in +    d; in < e; ++in, data += stride)
    Iter::iter4(*in, data, *accum);
}


template<typename Iter>
inline void iterateSimple(Layout cl, Layout pl, Layout wl, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  if (!cl) return;
  assert(pl);
  assert(wl);
  assert(c_neurons);
  assert(p_neurons);
  assert(weights);
  assert(cl.isSubLayoutOf(wl));

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

  int p_h    = pl.getH();
  int p_w    = pl.getW();
  int p_d    = pl.getD();
  int p_sxz  = pl.sx*pl.sz;
  int p_swz  = p_w*pl.sz;
  int p_shxz = p_h*p_sxz;
  int p_dy   = p_sxz - p_swz;
  int p_dx   = pl.sz - p_d;

  int w_w    = wl.getW();
  int w_d    = wl.getD();
  int w_dz   = p_h*p_w*p_d;
  int w_dx   = (w_d - c_d)*w_dz;
  int w_dy   = (w_w - c_w)*w_d*w_dz;

  Neuron *icn = c_neurons + (cl.y0*c_sxz + cl.x0*cl.sz + cl.z0);
  p_neurons += pl.y0*p_sxz + pl.x0*pl.sz + pl.z0;

  Weight *iw = weights + (((cl.y0 - wl.y0)*w_w + cl.x0 - wl.x0)*w_d + cl.z0 - wl.z0)*w_dz;

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

    Neuron *ipn = p_neurons;
    for(Neuron *e = ipn + p_shxz; ipn < e; ipn += p_dy)
    for(Neuron *e = ipn +  p_swz; ipn < e; ipn += p_dx)
    for(Neuron *e = ipn +    p_d; ipn < e; ++ipn, ++iw)
      Iter::iter(*ipn, *iw, a);

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


template<typename Iter>
void iterateSimpleInv(Layout cl, Layout pl, Layout wl, Neuron *c_neurons, Neuron *p_neurons, Weight *weights) {
  if (!cl) return;
  assert(pl);
  assert(wl);
  assert(c_neurons);
  assert(p_neurons);
  assert(weights);
  assert(cl.isSubLayoutOf(wl));

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

  int p_h    = pl.getH();
  int p_w    = pl.getW();
  int p_d    = pl.getD();
  int p_sxz  = pl.sx*pl.sz;
  int p_swz  = p_w*pl.sz;
  int p_shxz = p_h*p_sxz;
  int p_dy   = p_sxz - p_swz;
  int p_dx   = pl.sz - p_d;

  int w_w    = wl.getW();
  int w_h    = wl.getH();
  int w_d    = wl.getD();
  int w_ddz  = w_h*w_w*w_d;
  int w_dx   = w_d - c_d;
  int w_dy   = (w_w - c_w)*w_d;

  Neuron *icn = c_neurons + (cl.y0*c_sxz + cl.x0*cl.sz + cl.z0);
  p_neurons += pl.y0*p_sxz + pl.x0*pl.sz + pl.z0;

  Weight *iw = weights + ((cl.y0 - wl.y0)*w_w + cl.x0 - wl.x0)*w_d + cl.z0 - wl.z0;

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

    Weight *iiw = iw;
    Neuron *ipn = p_neurons;
    for(Neuron *e = ipn + p_shxz; ipn < e; ipn += p_dy)
    for(Neuron *e = ipn +  p_swz; ipn < e; ipn += p_dx)
    for(Neuron *e = ipn +    p_d; ipn < e; ++ipn, iiw += w_ddz)
      Iter::iter(*ipn, *iiw, a);

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


Layout optimizeLayoutSimple(const Layout &layout) {
  Layout l = layout;
  if (l.x0 == 0 && l.x1 == l.sx)
    { l.x0 = l.y0*l.sx; l.x1 *= l.y1; l.sx *= l.sy; l.y0 = 0; l.y1 = l.sy = 1; }
  if (l.z0 == 0 && l.z1 == l.sz)
    { l.z0 = l.x0*l.sz; l.z1 *= l.x1; l.sz *= l.sx; l.x0 = 0; l.x1 = l.sx = 1; }
  return l;
}


template<Func func>
class LayerSimple: public Layer {
public:
  Layout optLayout;
  Layout prevOptLayout;
  Layout::List mtOptLayouts;
  Layout::List mtPrevOptLayouts;


  LayerSimple(Layer &prev, const Layout &layout, Weight *weights = nullptr):
    Layer(&prev, layout, layout.getActiveCount() * prev.back().layout.getActiveCount(), weights),
    optLayout(optimizeLayoutSimple(layout)),
    prevOptLayout(optimizeLayoutSimple(this->prev->layout))
  {
    if (ownWeights) fillWeights(-1, 1);
  }


  void split(int threadsCount) override {
    Layer::split(threadsCount);
    optLayout.split(mtOptLayouts, threadsCount);
    prevOptLayout.split(mtPrevOptLayouts, threadsCount);
  }


  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); }
    };
    iterateSimple<I>(mtOptLayouts[barrier.tid], prevOptLayout, optLayout, 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; }
    };
    iterateSimple<I>(mtOptLayouts[barrier.tid], prevOptLayout, optLayout, neurons, prev->neurons, weights);
  }
  
  
  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; }
    };
    iterateSimpleInv<I>(mtPrevOptLayouts[barrier.tid], optLayout, prevOptLayout, prev->neurons, neurons, weights);
  }

  
  void testPass() override {
    Layout cl = layout;
    Layout pl = prev->layout;
    
    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) {
      AccumReal a = 0;
      Neuron &cn = neurons[ (cy*cl.sx + cx)*cl.sz + cz ];
      int wi = ((cy-cl.y0)*cl.getW() + cx-cl.x0)*cl.getD() + cz-cl.z0;

      for(int py = pl.y0; py < pl.y1; ++py)
      for(int px = pl.x0; px < pl.x1; ++px)
      for(int pz = pl.z0; pz < pl.z1; ++pz) {
        Neuron &pn = prev->neurons[ (py*pl.sx + px)*pl.sz + pz ];
        int wii = ((wi*pl.getH() + py-pl.y0)*pl.getW() + px-pl.x0)*pl.getD() + pz-pl.z0;
        Weight &w = weights[wii];
        a += pn.v * w.w;
      }
      
      func(cn, a);
    }
  }

  
  void testBackpass() override {
    Layout cl = layout;
    Layout pl = prev->layout;
  
    for(int i = 0; i < prev->neuronsCount; ++i)
      prev->neurons[i].a.v = 0;
    
    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) {
      Neuron &cn = neurons[ (cy*cl.sx + cx)*cl.sz + cz ];
      int wi = ((cy-cl.y0)*cl.getW() + cx-cl.x0)*cl.getD() + cz-cl.z0;

      for(int py = pl.y0; py < pl.y1; ++py)
      for(int px = pl.x0; px < pl.x1; ++px)
      for(int pz = pl.z0; pz < pl.z1; ++pz) {
        Neuron &pn = prev->neurons[ (py*pl.sx + px)*pl.sz + pz ];
        int wii = ((wi*pl.getH() + py-pl.y0)*pl.getW() + px-pl.x0)*pl.getD() + pz-pl.z0;
        Weight &w = weights[wii];
        
        pn.a.v += w.w * cn.d;
        w.w += pn.v * cn.d;
      }
    }

    for(int i = 0; i < prev->neuronsCount; ++i) {
      Neuron &pn = prev->neurons[i];
      pn.d *= pn.a.v;
      pn.a.v = 0;
    }
  }
};


#endif