#ifndef LAYER_CONV_INC_CPP

#define LAYER_CONV_INC_CPP

#include "tga.inc.cpp"

#include "font.inc.cpp"

#include "layer.simple.inc.cpp"

struct Kernel {

  int sx, sy;

  int dx, dy;

  int ox, oy;

  inline Kernel():

    sx(), sy(), dx(), dy(), ox(), oy() { }

  inline Kernel(int sx, int sy, int dx, int dy, int ox, int oy):

    sx(sx), sy(sy), dx(dx), dy(dy), ox(ox), oy(oy) { }

  inline Kernel(int s, int d, int o):

    sx(s), sy(s), dx(d), dy(d), ox(o), oy(o) { }

  inline operator bool() const

    { return sx > 0 && sy > 0 && dx > 0 && dy > 0; }

  void print(const char *prefix = nullptr) const {

    if (prefix && *prefix) printf("%s: ", prefix);

    printf("x(sdo): %d %d %d, y(sdo): %d %d %d\n", sx, dx, ox, sy, dy, oy);

  }

  void printYX(const char *prefix = nullptr) const {

    if (prefix && *prefix) printf("%s: ", prefix);

    printf("y(sdo): %d %d %d, x(sdo): %d %d %d\n", sy, dy, oy, sx, dx, ox);

  }

};

template<typename iter=""></typename>

void iterateTestConvolution(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 = ((cy - cl.y0)*cl.getW() + cx - cl.x0)*cl.getD() + cz - cl.z0;

      wi = ((wi*k.sy + 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=""></typename>

void iterateConvolution(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 k_syxd = k.sy*k_sxd;

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

  int p_ddx  = pl.sz - p_d;

  int w_w    = wl.getW();

  int w_d    = wl.getD();

  int w_dx   = (w_d - c_d)*k_syxd;

  int w_dy   = (w_w - c_w)*w_d*k_syxd;

  int cx0    = cl.x0 - wl.x0;

  int cy0    = cl.y0 - wl.y0;

  int cz0    = cl.z0 - wl.z0;

  Neuron *icn = c_neurons + (cl.y0*cl.sx + cl.x0)*cl.sz + cl.z0;

  Neuron *ipn = p_neurons + ((pl.y0 + cy0*k.dy + k.oy)*pl.sx + pl.x0 + cx0*k.dx + k.ox)*pl.sz + pl.z0;

  Weight *iw = weights + ((cy0*w_w + cx0)*w_d + cz0)*k_syxd;

  for(Neuron *e = icn + c_shxz; icn < e; icn += c_dy, ipn += p_dy, iw += w_dy)

  for(Neuron *e = icn +  c_swz; icn < e; icn += c_dx, ipn += p_dx, iw += w_dx)

  for(Neuron *e = icn +    c_d; icn < e; ++icn) {

    typename Iter::AccumType a;

    Iter::init(*icn, a);

    Neuron *iipn = ipn;

    for(Weight *e = iw + k_syxd; iw < e; 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=""></typename>

void iterateConvolutionPoint(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;

  int k_sxd  = k.sx*p_d;

  int k_syxd = k.sy*k_sxd;

  int w_w    = wl.getW();

  int w_d    = wl.getD();

  int w_dz   = k_syxd - p_d;

  int w_dx   = (w_d - c_d)*k_syxd;

  int w_dy   = (w_w - c_w)*w_d*k_syxd;

  int cx0    = cl.x0 - wl.x0;

  int cy0    = cl.y0 - wl.y0;

  int cz0    = cl.z0 - wl.z0;

  Neuron *icn = c_neurons + (cl.y0*cl.sx + cl.x0)*cl.sz + cl.z0;

  Neuron *ipn = p_neurons + ((pl.y0 + cy0*k.dy + k.oy + ky)*pl.sx + pl.x0 + cx0*k.dx + k.ox + kx)*pl.sz + pl.z0;

  Weight *iw = weights + ((cy0*w_w + cx0)*w_d + cz0)*k_syxd + ky*k_sxd + kx*p_d;

  for(Neuron *e = icn + c_shxz; icn < e; icn += c_dy, ipn += p_dy, iw += w_dy)

  for(Neuron *e = icn +  c_swz; icn < e; icn += c_dx, ipn += p_dx, iw += w_dx)

  for(Neuron *e = icn +    c_d; icn < e; ++icn,       ipn -= p_d,  iw += w_dz)

  for(Neuron *e = ipn +    p_d; ipn < e; ++ipn, ++iw)

    Iter::iter2(*icn, *ipn, *iw);

}

bool saveConvDemoImage(const char *filename, int count, int ksx, int ksy, int ksz, const Weight *weights) {

  int cols = count;

  int rows = ksz + 1;

  int w = 1 + cols*(ksx + 1);

  int h = 10 + rows*(ksy + 1);

  std::vector<unsigned char=""> pixels(w*h*3, 0);</unsigned>

  WeightReal range = 0;

  for(const Weight *iw = weights, *e = iw + count*ksx*ksy*ksz; iw < e; ++iw) {

    WeightReal r = fabs(iw->w);

    if (range < r) range = r;

  }

  const unsigned char white[] = { 255, 255, 255 };

  imgPrintf(pixels.data(), w, h, 3, 1, 1, white, "%f", range);

  // rgb row

  for(int i = 0; i < count; ++i)

  for(int ky = 0; ky < ksy; ++ky)

  for(int kx = 0; kx < ksx; ++kx) {

    int y0 = 10;

    int x0 = i*(ksx + 1) + 1;

    unsigned char *p = &pixels[ ((y0 + ky)*w + x0 + kx)*3 ];

    for(int kz = 0; kz < 3; ++kz) {

      if (kz < ksz) {

        WeightReal x = weights[ ((i*ksy + ky)*ksx + kx)*3 + kz ].w;

        x /= range;

        x = (x + 0.5)*256;

        unsigned char c = x < 0 ? 0 : x > 255 ? 255 : (unsigned char)x;

        p[kz] = c;

      } else {

        p[kz] = 0;

      }

    }

  }

  // gray rows

  for(int i = 0; i < count; ++i)

  for(int kz = 0; kz < ksz; ++kz)

  for(int ky = 0; ky < ksy; ++ky)

  for(int kx = 0; kx < ksx; ++kx) {

    WeightReal x = weights[ ((i*ksy + ky)*ksx + kx)*ksz + kz ].w;

    x /= range;

    x = (x + 0.5)*256;

    unsigned char c = x < 0 ? 0 : x > 255 ? 255 : (unsigned char)x;

    int y0 = (kz + 1)*(ksy + 1) + 10;

    int x0 = i*(ksx + 1) + 1;

    unsigned char *p = &pixels[ ((y0 + ky)*w + x0 + kx)*3 ];

    p[0] = p[1] = p[2] = c;

    //if (c ==   0) p[0] = p[1] = 0; // blue un underflow

    //if (c == 255) p[1] = p[2] = 0; // red on overflow

  }

  std::string fn(filename);

  fn += ".tga";

  return tgaSave(fn.c_str(), pixels.data(), w, h, 3);

}

template<func func=""></func>

class LayerConv: public Layer {

public:

  Kernel kernel;

  LayerConv(Layer &prev, const Layout &layout, const Kernel &kernel, Weight *weights = nullptr):

    Layer(&prev, layout, layout.getActiveCount()*kernel.sx*kernel.sy*prev.back().layout.getD(), weights),

    kernel(kernel)

  {

    assert(kernel);

    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); }

    };

    iterateConvolution(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; }

    };

    iterateConvolution(mtLayouts[barrier.tid], prev->layout, layout, kernel, neurons, prev->neurons, weights);

  }

  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) {

      iterateConvolutionPoint(mtLayouts[barrier.tid], prev->layout, layout, kernel, kx, ky, neurons, prev->neurons, weights);

      barrier.wait();

    }

    iterateNeurons(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); }

    };

    iterateTestConvolution(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; w.w += a.v * n.v; }

      static inline void iter3(Neuron &n) { n.d *= n.a.v; n.a.v = 0; }

    };

    clearAccum();

    iterateTestConvolution(layout, prev->layout, kernel, neurons, prev->neurons, weights);

    iterateNeurons(prev->layout, prev->neurons);

    clearAccum();

  }

};

template<func func=""></func>

class LayerDeconv: public Layer {

public:

  Kernel kernel;

  LayerDeconv(Layer &prev, const Layout &layout, const Kernel &kernel, Weight *weights = nullptr):

    Layer(&prev, layout, prev.back().layout.getActiveCount()*kernel.sx*kernel.sy*layout.getD(), weights),

    kernel(kernel)

  {

    assert(kernel);

    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) {

      iterateConvolutionPoint(mtPrevLayouts[barrier.tid], layout, prev->layout, kernel, kx, ky, prev->neurons, neurons, weights);

      barrier.wait();

    }

    iterateNeurons(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; }

    };

    iterateConvolution(mtPrevLayouts[barrier.tid], layout, prev->layout, kernel, prev->neurons, 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; }

    };

    iterateConvolution(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();

    iterateTestConvolution(prev->layout, layout, kernel, prev->neurons, neurons, weights);

    iterateNeurons(layout, neurons);

    clearAccum();

  }

  void testBackpass() override {

    struct I: public Iter {

      struct AccumType: public Accum { NeuronReal vv; };

      static inline void init(Neuron &n, AccumType &a) { a.v = 0; a.vv = n.v; }

      static inline void iter(Neuron &n, Weight &w, AccumType &a) { a.v += n.d * w.w; w.w += n.d * a.vv; }

      static inline void done(Neuron &n, AccumType &a) { n.d *= a.v; }

    };

    iterateTestConvolution(prev->layout, layout, kernel, prev->neurons, neurons, weights);

  }

};

#endif