#ifndef NNLAYER3_INC_CPP
#define NNLAYER3_INC_CPP


#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <cassert>

#include <vector>
#include <algorithm>


typedef float Real;


template<typename T>
bool twrite(const T &x, FILE *f)
  { return fwrite(&x, sizeof(T), 1, f); }
template<typename T>
bool tread(T &x, FILE *f)
  { return fread(&x, sizeof(T), 1, f); }


struct Neuron {
  Real v, d;
};


class Layer {
public:
  Layer *prev, *next;

  int sx, sy, sz, sl;
  Neuron *neurons;
  Real *weights;
  int *invWeights;
  int *invNeurons;


  explicit Layer(Layer *prev = nullptr, int sx = 0, int sy = 0, int sz = 0, int sl = 0):
    prev(), next(), sx(), sy(), sz(), sl(), neurons(), weights(), invWeights(), invNeurons()
  {
    while(prev && prev->next) prev = prev->next;
    if (prev) prev->next = this;
    this->prev = prev;
    init(sx, sy, sz, sl);
  }


  virtual ~Layer() {
    clear();
    if (next) delete next;
    if (prev) prev->next = nullptr;
  }


  void clear() {
    if (next) next->clear();
    if (neurons) delete[] neurons;
    if (weights) delete[] weights;
    if (invWeights) delete[] invWeights;
    sx = sy = sz = sl = 0;
    neurons = nullptr;
    weights = nullptr;
    invWeights = nullptr;
    invNeurons = nullptr;
  }


  bool init(int sx, int sy, int sz, int sl = 0) {
    clear();
    if (prev && !prev->countNeurons()) return false;
    if (sx <= 0 || sy <= 0 || sz <= 0) return false;
    if (prev ? (sl <= 0 || sl > prev->sx || sl > prev->sy) : sl != 0) return false;
    printf("init %d %d %d %d\n", sx, sy, sz, sl);

    this->sx = sx;
    this->sy = sy;
    this->sz = sz;

    int size = sx*sy*sz;
    neurons = new Neuron[size];
    memset(neurons, 0, sizeof(*neurons)*size);

    if (prev) {
      int psx = prev->sx;
      int psy = prev->sy;
      int psz = prev->sz;
      int psize = psx*psy*psz;

      this->sl = sl;
      int wsize = size*sl*sl*psz;
      weights = new Real[wsize];
      double k = 1.0/(sl*sl*psz);
      for(int i = 0; i < wsize; ++i)
        weights[i] = Real( (rand()/(double)RAND_MAX*2 - 1)*k );

      struct Link {
        int n, w;
        inline bool operator< (const Link &b) const
          { return n < b.n ? true : (b.n < n ? false : w < b.w); }
      } *links = new Link[wsize], *il = links;

      int dx = sx > 1 ? sx - 1 : 1;
      int dy = sy > 1 ? sy - 1 : 1;

      for(int y = 0; y < sy; ++y)
      for(int x = 0; x < sx; ++x) {
        int py = y*(psy-sl)/dy;
        int px = x*(psx-sl)/dx;
        assert(py >= 0 && py <= psy-sl);
        assert(px >= 0 && px <= psx-sl);

        for(int z = 0; z < sz; ++z)
        for(int ly = 0; ly < sl; ++ly)
        for(int lx = 0; lx < sl; ++lx)
        for(int pz = 0; pz < psz; ++pz, ++il) {
          assert(il < links + wsize);
          il->n = ((py+ly)*psx + px+lx)*psz + pz;
          il->w = il - links;
          assert(il->n >= 0 && il->n < psize);
        }
      }
      assert(il == links + wsize);
      std::sort(links, il);

      invWeights = new int[wsize + psize + 1];
      invNeurons = invWeights + wsize;
      invNeurons[0] = 0;
      int pni = 0;
      for(int i = 0; i < wsize; ++i) {
        assert(pni == links[i].n || pni == links[i].n + 1);
        invWeights[i] = links[i].w;
        if (pni == links[i].n)
          invNeurons[pni++] = i;
      }
      assert(pni == psize);
      invNeurons[psize] = wsize;

      delete[] links;
    }
    return true;
  }

  bool save(FILE *f) const {
    return twrite(sx, f) && twrite(sy, f) && twrite(sz, f) && twrite(sl, f)
        && (!weights || fwrite(weights, sizeof(*weights)*sx*sy*sz*sl*sl*prev->sz, 1, f));
  }

  bool load(FILE *f) {
    clear();
    int sx = 0, sy = 0, sz = 0, sl = 0;
    return tread(sx, f) && tread(sy, f) && tread(sz, f) && tread(sl, f) && init(sx, sy, sz, sl)
        && (!weights || fread(weights, sizeof(*weights)*sx*sy*sz*sl*sl*prev->sz, 1, f));
  }

  bool saveAll(FILE *f) const
    { return save(f) && (!next || next->saveAll(f)); }
  bool loadAll(FILE *f)
    { return load(f) && (!next || next->loadAll(f)); }

  bool saveAll(const char *filename) const {
    assert(!prev);
    FILE *f = fopen(filename, "wb");
    if (!f)
      { printf("cannot open file for write: %s\n", filename); return false; }
    int count = totalLayers();
    if (!twrite(count, f) || !saveAll(f))
      { printf("cannot save to file: %s\n", filename); fclose(f); return false; }
    fclose(f);
    return true;
  }

  bool loadAll(const char *filename) {
    assert(!prev);
    FILE *f = fopen(filename, "rb");
    if (!f)
      { printf("cannot open file for read: %s\n", filename); return false; }
    int count = 0;
    if (!tread(count, f) || count <= 0)
      { printf("cannot load from file: %s\n", filename); fclose(f); return false; }
    if (next)
      delete next;
    while(--count)
      new Layer(this);
    if (!loadAll(f))
      { printf("cannot load from file: %s\n", filename); fclose(f); return false; }
    fclose(f);
    return true;
  }

  inline int countNeurons() const
    { return sx*sy*sz; }
  inline int linksPerNeuron() const
    { return prev ? sl*sl*prev->sz : 0; }
  inline int countLinks() const
    { return countNeurons() * linksPerNeuron(); }
  inline size_t memSize() const {
    return sizeof(*neurons)*countNeurons()
         + (sizeof(*weights) + sizeof(*invWeights))*countLinks()
         + (prev ? sizeof(*invNeurons)*(prev->countNeurons() + 1) : 0);
  }

  int totalLayers() const
    { int t = 0; for(const Layer *l = this; l; l = l->next) ++t; return t; }
  int totalNeurons() const
    { int t = 0; for(const Layer *l = this; l; l = l->next) t += l->countNeurons(); return t; }
  int totalLinks() const
    { int t = 0; for(const Layer *l = this; l; l = l->next) t += l->countLinks(); return t; }
  size_t totalMemSize() const
    { size_t t = 0; for(const Layer *l = this; l; l = l->next) t += l->memSize(); return t; }

  inline Layer& front()
    { Layer *l = this; while(l->prev) l = l->prev; return *l; }
  inline Layer& back()
    { Layer *l = this; while(l->next) l = l->next; return *l; }


  void pass(int y0, int y1) {
    if (!prev) return;

    int sx = this->sx;
    int sy = this->sy;
    int sz = this->sz;
    int sxz = sx*sz;

    int psx = prev->sx;
    int psy = prev->sy;
    int psz = prev->sz;
    int psxz = psx*psz;
    Neuron *pneurons = prev->neurons;

    int sl = this->sl;
    int slpz = sl*psz;
    int sllpz = sl*slpz;

    int ldy = psxz - slpz;

    int kpx = psx-sl, dpx = sx>1 ? sx-1 : 1;
    int kpy = psy-sl, dpy = sy>1 ? sy-1 : 1;

    Real *iw = weights + y0*sxz*sllpz;
    Neuron *in = neurons + y0*sxz;
    int fpy = y0*kpy;
    for(Neuron *e = in + (y1-y0)*sxz; in < e; fpy += kpy) {
      Neuron *pnrow = pneurons + fpy/dpy*psxz;

      int fpx = 0;
      for(Neuron *e = in + sxz; in < e; fpx += kpx) {
        Neuron *pn = pnrow + fpx/dpx*psz;

        for(Neuron *e = in + sz; in < e; ++in) {
          double s = 0;

          Neuron *ipn = pn;
          for(Real *e = iw + sllpz; iw < e; ipn += ldy)
          for(Real *e = iw + slpz; iw < e; ++iw, ++ipn)
            s += *iw * ipn->v;

          // exp sigmoid
          //double ss = 1/(1 + exp(-s));
          //in->v = ss;
          //in->d = ss * (1-ss);

          // 1/(x+1) sigmoid
          double ss = 1/(1+fabs(s));
          double ss2 = ss*0.5;
          in->v = s > 0 ? 1 - ss2 : ss2;
          in->d = ss2 * ss;
        }
      }
    }
  }


  void backpassWeights(int y0, int y1) {
    assert(prev);

    int sx = this->sx;
    int sy = this->sy;
    int sz = this->sz;
    int sxz = sx*sz;

    int psx = prev->sx;
    int psy = prev->sy;
    int psz = prev->sz;
    int psxz = psx*psz;
    Neuron *pneurons = prev->neurons;

    int sl = this->sl;
    int slpz = sl*psz;
    int sllpz = sl*slpz;

    int ldy = psxz - slpz;

    int kpx = psx-sl, dpx = sx>1 ? sx-1 : 1;
    int kpy = psy-sl, dpy = sy>1 ? sy-1 : 1;

    Real *iw = weights + y0*sxz*sllpz;
    Neuron *in = neurons + y0*sxz;
    int fpy = y0*kpy;
    for(Neuron *e = in + (y1-y0)*sxz; in < e; fpy += kpy) {
      Neuron *pnrow = pneurons + fpy/dpy*psxz;

      int fpx = 0;
      for(Neuron *e = in + sxz; in < e; fpx += kpx) {
        Neuron *pn = pnrow + fpx/dpx*psz;

        for(Neuron *e = in + sz; in < e; ++in) {
          Real d = in->d;

          Neuron *ipn = pn;
          for(Real *e = iw + sllpz; iw < e; ipn += ldy)
          for(Real *e = iw + slpz; iw < e; ++iw, ++ipn)
            *iw += ipn->v * d;
        }
      }
    }
  }

  template<bool WithWeights>
  void backpassTpl(int y0, int y1) {
    assert(next);

    Neuron *nneurons = next->neurons;
    Real *nweights = next->weights;
    int *nInvWeights = next->invWeights;
    int lpn = next->linksPerNeuron();

    int sxz = sx*sz;

    Neuron *in = neurons + y0*sxz;
    int *inni = next->invNeurons + y0*sxz;
    for(Neuron *e = in + (y1-y0)*sxz; in < e; ++in) {
      Real v = in->v;
      double s = 0;

      int *iw = nInvWeights + *inni++;
      for(int *e = nInvWeights + *inni; iw < e; ++iw) {
        int wi = *iw;
        Real d = nneurons[wi/lpn].d;
        Real &nw = nweights[wi];
        s += nw * d;
        if (WithWeights) nw += d * v;
      }

      in->d *= s;
    }
  }

  void passAll() {
    if (prev) pass(0, sy);
    if (next) next->passAll();
  }

  void backpassAll() {
    if (!prev)
      return;
    if (!prev->prev)
      { backpassWeights(0, sy); return; }
    if (next)
      backpassTpl<true>(0, sy);
    return prev->backpassAll();
  }
};


#endif