#ifndef SEGMENT_CX4_INC_CPP

#define SEGMENT_CX4_INC_CPP

#include "segment.inc.cpp"

#include "func.inc.cpp"

#include "layer.conv.inc.cpp"

class SegmentCx4: public Segment {

public:

  enum {

    KSX = 4,

    KSY = 4,

    SX = 12,

    SY = 12,

    MSX = 5,

    MSY = 5,

  };

  const int msx, msy, msz;

  Neuron *m_neurons;

  Neuron *b_neurons;

  SegmentCx4(int sz, int msz, Weight *weights = nullptr):

    Segment(SX, SY, sz, msz*KSY*KSX*sz, weights), msx(MSX), msy(MSY), msz(msz)

  {

      m_neurons = new Neuron[msx*msy*msz + sx*sy*sz];

      b_neurons = m_neurons + msx*msy*msz;

      clear();

  }

  ~SegmentCx4()

    { delete[] m_neurons; }  

  void clear() override

    { memset(m_neurons, 0, sizeof(*m_neurons)*(msx*msy*msz + sx*sy*sz)); }

  inline void check(int x, int y, int z) {

    Segment::check(x, y, z);

    assert(layout.getD() == sz);

  }

  Quality pass(Barrier &barrier, int x, int y, int z, NeuronReal trainRatio) override {

    check(x, y, z);

    Layout l = layout;

    const int ksx = 4, ksy = 4;

    int tid = barrier.tid;

    int threads = barrier.threads;

    int sx = this->sx;

    int sy = this->sy;

    int sz = this->sz;

    int msx = this->msx;

    int msy = this->msy;

    int msz = this->msz;

    int ksxyz = ksx*ksy*sz;

    int fv_dkx = l.sz - sz;

    int fv_dky = (l.sx - ksx)*l.sz;

    NeuronReal *f_values = this->f_values + (y*l.sx + x)*l.sz + z;

    // stage 1: pass from front to mid

    Weight *iw = weights + tid*ksxyz;

    Neuron *imn = m_neurons + tid;

    NeuronReal *ifv = f_values;

    for(int mz = tid; mz < msz; mz += threads, iw += threads*ksxyz, imn += threads - msx*msy*msz, ifv = f_values)

    for(int my = 0; my < MSY; ++my, ifv += 2*(l.sx - MSX)*l.sz)

    for(int mx = 0; mx < MSX; ++mx, imn += msz, ifv += 2*l.sz) {

      AccumReal a = 0;

      Weight *iiw = iw;

      NeuronReal *iifv = ifv;

      for(int ky = 0; ky < KSY; ++ky, iifv += fv_dky)

      for(int kx = 0; kx < KSX; ++kx, iifv += fv_dkx)

      for(Weight *e = iiw + sz; iiw < e; ++iiw, ++iifv)

        a += *iifv * iiw->w;

      if (a > 0) imn->v = a, imn->d = 1; else imn->v = imn->d = 0;

    }

    barrier.wait();

    // stage 2: pass from mid to back and verify

    AccumReal qa = 0;

    for(int by = 2 + tid; by < 10; by += threads)

    for(int bx = 2; bx < 10; ++bx)

    for(int bz = 0; bz < sz; ++bz) {

      AccumReal a = 0;

      Neuron &bn = b_neurons[ (by*sx + bx)*sz + bz ];

      for(int ky = by%2; ky < ksy; ky += 2)

      for(int kx = bx%2; kx < ksx; kx += 2) {

        int mx = (bx - kx)/2;

        int my = (by - ky)/2;

        assert(mx >= 0 && mx < msx && (bx - kx)%2 == 0);

        assert(my >= 0 && my < msy && (by - ky)%2 == 0);

        for(int mz = 0; mz < msz; ++mz) {

          Neuron &mn = m_neurons[ (my*msx + mx)*msz + mz ];

          Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + bz ];

          a += mn.v * w.w;

        }

      }

      if (a > 0) bn.v = a, bn.d = 1; else bn.v = bn.d = 0;

      NeuronReal fn = f_values[ (by*l.sx + bx)*l.sz + bz ];

      NeuronReal d = fn - bn.v;

      bn.d *= d*trainRatio;

      qa += d*d;

    }

    Quality q(qa/(64*sz));

    if (trainRatio <= 0) return q;

    barrier.wait();

    // stage 3: backpass deltas

    for(int mz = tid; mz < msz; mz += threads)

    for(int my = 1; my < 4; ++my)

    for(int mx = 1; mx < 4; ++mx) {

      AccumReal a = 0;

      Neuron &mn = m_neurons[ (my*msx + mx)*msz + mz ];

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

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

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

        int bx = mx*2 + kx;

        int by = my*2 + ky;

        Neuron &bn = b_neurons[ (by*sx + bx)*sz + kz ];

        Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + kz ];

        a += bn.d * w.w;

      }

      mn.d *= a;

    }

    barrier.wait();

    // stage 4: update weights

    for(int mz = tid; mz < msz; mz += threads)

    for(int by = 4; by <  8; ++by)

    for(int bx = 4; bx <  8; ++bx)

    for(int bz = 0; bz < sz; ++bz) {

      Neuron &bn = b_neurons[ (by*sx + bx)*sz + bz ];

      NeuronReal fv = f_values[ (by*l.sx + bx)*l.sz + bz ];

      for(int ky = by%2; ky < ksy; ky += 2)

      for(int kx = bx%2; kx < ksx; kx += 2) {

        int mx = (bx - kx)/2;

        int my = (by - ky)/2;

        assert(mx >= 1 && mx < 4 && (bx - kx)%2 == 0);

        assert(my >= 1 && my < 4 && (by - ky)%2 == 0);

        Neuron &mn = m_neurons[ (my*msx + mx)*msz + mz ];

        Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + bz ];

        w.w += bn.d*mn.v + mn.d*fv;

      }

    }

    return q;

  }

  Quality testPass(int x, int y, int z, NeuronReal trainRatio) override {

    check(x, y, z);

    Layout l = layout;

    const int ksx = 4, ksy = 4;

    // stage 1: pass

    clear();

    for(int my = 0; my < msy; ++my)

    for(int mx = 0; mx < msx; ++mx)

    for(int mz = 0; mz < msz; ++mz) {

      AccumReal a = 0;

      Neuron &mn = m_neurons[ (my*msx + mx)*msz + mz ];

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

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

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

        int fx = x + mx*2 + kx;

        int fy = y + my*2 + ky;

        int fz = z + kz;

        NeuronReal fv = f_values[ (fy*l.sx + fx)*l.sz + fz ];

        Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + kz ];

        a += fv * w.w;

      }

      if (a < 0) { mn.v = mn.d = 0; continue; }

      mn.v = a; mn.d = 1;

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

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

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

        int bx = mx*2 + kx;

        int by = my*2 + ky;

        int bz = kz;

        Neuron &bn = b_neurons[ (by*sx + bx)*sz + bz ];

        Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + kz ];

        bn.a.v += a * w.w;

      }

    }

    // stage 2: finalize values and verify

    AccumReal qa = 0;

    for(int by = 2; by < 10; ++by)

    for(int bx = 2; bx < 10; ++bx)

    for(int bz = 0; bz < sz; ++bz) {

        Neuron &bn = b_neurons[ (by*sx + bx)*sz + bz ];

        if (bn.a.v > 0) bn.v = bn.a.v, bn.d = 1; else bn.v = bn.d = 0;

        bn.a.v = 0;

        NeuronReal fn = f_values[ ((y + by)*l.sx + x + bx)*l.sz + z + bz ];

        NeuronReal d = fn - bn.v;

        bn.d *= d*trainRatio;

        qa += d*d;

    }

    Quality q(qa/(64*sz));

    if (trainRatio <= 0) return q;

    // stage 3: backpass deltas

    for(int my = 0; my < msy; ++my)

    for(int mx = 0; mx < msx; ++mx)

    for(int mz = 0; mz < msz; ++mz) {

      AccumReal a = 0;

      Neuron &mn = m_neurons[ (my*msx + mx)*msz + mz ];

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

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

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

        int bx = mx*2 + kx;

        int by = my*2 + ky;

        int bz = kz;

        Neuron &bn = b_neurons[ (by*sx + bx)*sz + bz ];

        Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + kz ];

        a += bn.d * w.w;

      }

      mn.d *= a;

    }

    // stage 4: update weights

    for(int by = 4; by <  8; ++by)

    for(int bx = 4; bx <  8; ++bx)

    for(int bz = 0; bz < sz; ++bz) {

      Neuron &bn = b_neurons[ (by*sx + bx)*sz + bz ];

      NeuronReal fv = f_values[ ((y + by)*l.sx + x + bx)*l.sz + z + bz ];

      for(int ky = by%2; ky < ksy; ky += 2)

      for(int kx = bx%2; kx < ksx; kx += 2)

      for(int mz = 0; mz < msz; ++mz) {

        int mx = (bx - kx)/2;

        int my = (by - ky)/2;

        assert(mx >= 1 && mx < 4 && (bx - kx)%2 == 0);

        assert(my >= 1 && my < 4 && (by - ky)%2 == 0);

        Neuron &mn = m_neurons[ (my*msx + mx)*msz + mz ];

        Weight &w = weights[ ((mz*ksy + ky)*ksx + kx)*sz + bz ];

        w.w += bn.d*mn.v + mn.d*fv;

      }

    }

    return q;

  }

  bool saveDemo() override

    { return !filename || saveConvDemoImage(filename, msz, 4, 4, sz, weights); }

};

#endif