#ifndef LAYER_CONV_TEST_INC_CPP

#define LAYER_CONV_TEST_INC_CPP

#include "layer.test.inc.cpp"

#include "layer.conv.inc.cpp"

#include "layer.conv.shared.inc.cpp"

class ConvTest: public LayerTest {

public:

  static void init(const Layout &cl, const Layout &pl, const Kernel &k, bool shared = false)

    { Test::init(cl.getCount(), pl.getCount(), (shared ? 1 : cl.getActiveCount())*k.sx*k.sy*pl.getD()); }

  static bool verifyWeights(const char *name, const Layout &cl, const Layout &pl, const Kernel &k) {

    Stage st(name);

    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;

      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;

        int px = pl.x0 + (cx - cl.x0)*k.dx + k.ox + kx;

        int py = pl.y0 + (cy - cl.y0)*k.dy + k.oy + ky;

        if ( px < pl.x0 || px >= pl.x1

          || py < pl.y0 || py >= pl.y1 ) continue;

        int pi = (py*pl.sx + px)*pl.sz + pz;

        int s = (int)p_neurons.size();

        int w = weights[wi].i;

        int i = ci*s + pi + 1;

        if (w != i) {

          int ww = w;

          int wpz = ww%pl.sz; ww /= pl.sz;

          int wpx = ww%pl.sx; ww /= pl.sx;

          int wpy = ww%pl.sy; ww /= pl.sy;

          int wcz = ww%cl.sz; ww /= cl.sz;

          int wcx = ww%cl.sx; ww /= cl.sx;

          int wcy = ww;

          printf(

            "wrong index: %d = ((%d*%d + %d)*%d + %d)*%d + (%d*%d + %d)*%d + %d + 1,\n"

            "expected:    %d = ((%d*%d + %d)*%d + %d)*%d + (%d*%d + %d)*%d + %d + 1\n"

            "wi %d, ky %d, kx %d \n",

            w,

            wcy, cl.sx, wcx, cl.sz, wcz, s,

            wpy, pl.sx, wpx, pl.sz, wpz,

            i,

            cy, cl.sx, cx, cl.sz, cz, s,

            py, pl.sx, px, pl.sz, pz,

            wi, ky, kx );

          pl.printYXZ("prev layout");

          cl.printYXZ("curr layout");

          k.printYX("kernel");

          ++errors;

          return st;

        }

      }

    }

    return st;

  }

  static bool testIterators(const char *name, const Layout &cl, const Layout &pl, const Kernel &k, int threads) {

    Stage st(name);

    assert(cl && pl && k && threads > 0);

    Layout::List clist, plist;

    cl.split(clist, threads);

    pl.split(plist, threads);

    struct I: public Iter {

      static inline void init(Neuron &n, Iter::AccumType &a) { ++n.a.i; a.i = (AccumInt)(&n - c_neurons.data()); }

      static inline void iter(Neuron &n, Weight &w, Iter::AccumType &a) {

        if (w.i)

          ++errors;

        w.i = (WeightInt)(&n - p_neurons.data() + a.i*p_neurons.size() + 1);

      }

      static inline void iter2(Neuron &cn, Neuron &pn, Weight &w) {

        if (w.i)

          ++errors;

        w.i = (WeightInt)((&cn - c_neurons.data())*p_neurons.size() + &pn - p_neurons.data() + 1);

        pn.v = pn.v + 1;

      }

    };

    if (threads == 1) {

      Stage st("iterateTestConvolution");

      init(cl, pl, k);

      iterateTestConvolution(cl, pl, k, c_neurons.data(), p_neurons.data(), weights.data());

      verifyNeurons("conv-neurons", cl, c_neurons.data());

      verifyWeights("conv-weights", cl, pl, k);

    }

    {

      Stage st("iterateConvolution");

      init(cl, pl, k);

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

        iterateConvolution(clist[i], pl, cl, k, c_neurons.data(), p_neurons.data(), weights.data());

      verifyNeurons("conv-neurons", cl, c_neurons.data());

      verifyWeights("conv-weights", cl, pl, k);

    }

    {

      Stage st("iterateConvolutionPoint");

      init(cl, pl, k);

      int e = errors;

      for(int ky = 0; ky < k.sy && errors == e; ++ky)

      for(int kx = 0; kx < k.sx && errors == e; ++kx) {

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

          iterateConvolutionPoint(clist[i], pl, cl, k, kx, ky, c_neurons.data(), p_neurons.data(), weights.data());

        if (!verifyNeuronsAccum(pl, p_neurons.data(), cl.getD(), true))

          printf("kx: %d, ky: %d\n", kx, ky), k.printYX("kernel");

      }

      verifyNeurons("conv-neurons", pl, p_neurons.data(), true);

      verifyWeights("conv-weights", cl, pl, k);

    }

    return st;

  }

  static bool testIterators(const char *name, const Layout &cl, const Layout &pl, const Kernel &k) {

    Stage st(name);

    testIterators( "single-thread", cl, pl, k,   1 );

    testIterators( "2-threads",     cl, pl, k,   2 );

    testIterators( "7-threads",     cl, pl, k,   7 );

    testIterators( "8-threads",     cl, pl, k,   8 );

    testIterators( "512-threads",   cl, pl, k, 512 );

    return st;

  }

  static bool test(const char *name, const Layout &cl, const Layout &pl, const Kernel &k) {

    Stage st(name);

    testIterators("iterators", cl, pl, k);

    {

      Layer l(nullptr, pl);

      new LayerConv<funcSigmoidExp>(l, cl, k);

      testLayer("LayerConv", l);

    }

    {

      Layer l(nullptr, cl);

      new LayerDeconv<funcSigmoidExp>(l, pl, k);

      testLayer("LayerDeconv", l);

    }

    {

      Layer l(nullptr, pl);

      new LayerConvShared<funcSigmoidExp>(l, cl, k);

      testLayer("LayerConvShared", l);

    }

    {

      Layer l(nullptr, cl);

      new LayerDeconvShared<funcSigmoidExp>(l, pl, k);

      testLayer("LayerDeconvShared", l);

    }

    return st;

  }

  static bool test(const char *name = "convolution") {

    Stage st(name);

    test( "square", Layout(64, 64, 4), Layout(128, 128, 4).expandXY(2),                 Kernel(5, 2, -2)           );

    test( "rect1",  Layout(63, 43, 5), Layout( 63,  85, 3).expandX(2).   expandY(3),    Kernel(5, 7, 1, 2, -2, -3) );

    test( "rect2",  Layout(43, 63, 3), Layout( 85,  63, 5).expandX(3).   expandY(3, 2), Kernel(7, 5, 2, 1, -3, -2) );

    test( "rect3",  Layout(64, 48, 5), Layout( 64,  96, 3).expandX(1, 2).expandY(3, 1), Kernel(4, 6, 1, 2, -1, -3) );

    test( "pad",    Layout(64, 48, 5).expandX(3, 4).expandY(4, 3).expandZ(5, 4),

                    Layout(64, 96, 3).expandX(6, 5).expandY(7, 6).expandZ(0, 1),        Kernel(4, 6, 1, 2, -1, -3) );

    return st;

  }

};

#endif