#ifndef LAYER_TEST_INC_CPP
#define LAYER_TEST_INC_CPP


#include <thread>

#include "layer.inc.cpp"



class Test {
public:
  class Stage {
  public:
    const int errors;
    inline explicit Stage(const char *name): errors(Test::errors) {
      for(int i = 0; i < level; ++i) printf("- ");
      printf("%s\n", name);
      fflush(stdout);
      ++level;
    }
    inline ~Stage() {
      --level;
      if (!*this) {
        for(int i = 0; i < level; ++i) printf("- ");
        printf("FAILED\n");
      }
      fflush(stdout);
    }
    operator bool() { return Test::errors == errors; }
  };

private:
  static int level;

protected:
  static std::vector<Neuron> c_neurons;
  static std::vector<Neuron> p_neurons;
  static std::vector<Weight> weights;

public:
  static int errors;


  static void init(int c_count, int p_count, int w_count) {
    Neuron n = {};
    Weight w = {};

    c_neurons.clear();
    p_neurons.clear();
    weights.clear();
    c_neurons.resize(c_count, n);
    p_neurons.resize(p_count, n);
    weights.resize(w_count, w);
  }


  static bool verifyNeurons(const char *name, const Layout &l, const Neuron *neurons, bool ignorePadded = false) {
    Stage st(name);
    for(int y = 0; y < l.sy; ++y)
    for(int x = 0; x < l.sx; ++x)
    for(int z = 0; z < l.sz; ++z) {
      int n = neurons[ (y*l.sx + x)*l.sz + z ].a.i;
      int i = x >= l.x0 && x < l.x1
           && y >= l.y0 && y < l.y1
           && z >= l.z0 && z < l.z1;
      if (ignorePadded ? i && n != i : n != i) {
        printf(
          "wrong neuron mark %d, expected %d (%d, %d, %d)\n",
          n, i, y, x, z );
        l.printYXZ("layout");
        ++errors;
        return st;
      }
    }
    return st;
  }
  
  
  static bool verifyNeuronIndices(const char *name, const Layout &l, const Neuron *neurons, int base = 1, int stride = 1) {
    Stage st(name);
    for(int y = 0; y < l.sy; ++y)
    for(int x = 0; x < l.sx; ++x)
    for(int z = 0; z < l.sz; ++z) {
      bool active = x >= l.x0 && x < l.x1
                 && y >= l.y0 && y < l.y1
                 && z >= l.z0 && z < l.z1;
      
      int n = neurons[ (y*l.sx + x)*l.sz + z ].a.i;
      int i = (((y - l.y0)*l.getW() + x - l.x0)*l.getD() + z - l.z0)*stride + base;
      
      if (!active) i = 0;
      
      if (n != i) {
        printf(
          "wrong neuron mark %d, expected %d (%d, %d, %d)\n",
          n, i, y, x, z );
        l.printYXZ("layout");
        ++errors;
        return st;
      }
    }
    return st;
  }
  
  
  static bool verifyNeuronsAccum(const Layout &l, Neuron *neurons, int accum = 1, bool ignoreBounds = false) {
    for(int y = 0; y < l.sy; ++y)
    for(int x = 0; x < l.sx; ++x)
    for(int z = 0; z < l.sz; ++z) {
      Neuron &n = neurons[ (y*l.sx + x)*l.sz + z ];
      int i = ( x >= l.x0 && x < l.x1
             && y >= l.y0 && y < l.y1
             && z >= l.z0 && z < l.z1 )*accum;
      if (ignoreBounds) i = accum;
      if (n.v != 0 && n.v != i) {
        printf(
          "wrong neuron mark %g, expected 0 or %d (%d, %d, %d)\n",
          n.v, i, y, x, z );
        l.printYXZ("layout");
        ++errors;
        return false;
      }
      if (n.v) n.a.i = 1;
      n.v = 0;
    }
    return true;
  }
  
  
  static bool testLayer(const char *name, Layer &l) {
    Stage st(name);

    assert(l.next);
    Layer &p = l;
    Layer &c = *l.next;
    
    
    struct H {
      Layer &p;
      Layer &c;

      std::vector<std::thread*> threads;
      std::atomic<unsigned int> counter;
      
      H(Layer &p, Layer &c): p(p), c(c), counter(0) { }
  
      void prepareData() {
        memcpy(c.neurons, c_neurons.data(), c.neuronsCount*sizeof(Neuron));
        memcpy(p.neurons, p_neurons.data(), p.neuronsCount*sizeof(Neuron));
        memcpy(c.weights, weights.data(), c.weightsCount*sizeof(Weight));
      }
  
      void applyDelta() {
        for(int i = 0; i < c.neuronsCount; ++i)
          c.neurons[i].d *= c_neurons[i].v - c.neurons[i].v;
      }
  
      void func(int tid) {
        Barrier barrier(counter, tid, threads.size());
        c.pass(barrier);
        barrier.wait();
        if (!tid) applyDelta();
        barrier.wait();
        c.backpassDeltas(barrier);
        barrier.wait();
        c.backpassWeights(barrier);
      }
      
      bool test(const char *name, int threadsCount) {
        Stage st(name);
        
        assert(threadsCount > 0);
        
        counter = 0;
        threads.clear();
        threads.resize(threadsCount, nullptr);
        
        prepareData();

        p.split(threadsCount);
        c.split(threadsCount);
        for(int i = 1; i < threadsCount; ++i) threads[i] = new std::thread(&H::func, this, i);
        func(0);
        for(int i = 1; i < threadsCount; ++i) { threads[i]->join(); delete threads[i]; }
        threads.clear();
        
        for(int i = 0; i < c.neuronsCount; ++i) {
          NeuronReal a = c.neurons[i].v;
          NeuronReal b = c_neurons[i + c.neuronsCount].v;
          if (fabs(a - b) > 1e-6)
            { printf("results differs at neuron %d, was %g, expected %g\n", i, a, b); ++errors; break; }
        }
        
        for(int i = 0; i < p.neuronsCount; ++i) {
          NeuronReal a = p.neurons[i].d;
          NeuronReal b = p_neurons[i + p.neuronsCount].d;
          if (fabs(a - b) > 1e-6)
            { printf("deltas differs at neuron %d, was %g, expected %g\n", i, a, b); ++errors; break; }
        }

        for(int i = 0; i < c.weightsCount; ++i) {
          WeightReal a = c.weights[i].w;
          WeightReal b = weights[i + c.weightsCount].w;
          if (fabs(a - b) > 1e-6)
            { printf("weights differs at %d, was %g, expected %g\n", i, a, b); ++errors; break; }
        }
        
        if (!st) {
          p.layout.printYXZ("prev layout");
          c.layout.printYXZ("curr layout");
        }
        
        return st;
      }
    } h(p, c);
    
    // make base data
    
    init(c.neuronsCount*2, p.neuronsCount*2, c.weightsCount*2);
    for(int i = 0; i < c.neuronsCount; ++i) c_neurons[i].v = rand()/(NeuronReal)RAND_MAX;
    for(int i = 0; i < p.neuronsCount; ++i) p_neurons[i].v = rand()/(NeuronReal)RAND_MAX;
    memcpy(weights.data(), c.weights, c.weightsCount*sizeof(Weight));
    
    h.prepareData();
    c.testPass();
    h.applyDelta();
    c.testBackpass();

    memcpy(&c_neurons[c.neuronsCount], c.neurons, c.neuronsCount*sizeof(Neuron));
    memcpy(&p_neurons[p.neuronsCount], p.neurons, p.neuronsCount*sizeof(Neuron));
    memcpy(&weights[c.weightsCount], c.weights, c.weightsCount*sizeof(Weight));

    h.test("single-thread", 1);
    h.test("2-threads", 2);
    h.test("7-threads", 7);
    h.test("8-threads", 8);
    //h.test("512-threads", 512);
    
    return st;
  }
};


int Test::level = 0;
std::vector<Neuron> Test::c_neurons;
std::vector<Neuron> Test::p_neurons;
std::vector<Weight> Test::weights;
int Test::errors = 0;



#endif