#include "segment.inc.cpp"
#include "segment.test.inc.cpp"
#include "segment.cx4.inc.cpp"


class BenchmarkSegment {
public:
  static std::vector<NeuronReal> values;
  static std::vector<Weight> weights;
  
  static void runSegment(const char *name, Segment &s, int threads, int repeats, bool backpass) {
    class H: public ThreadControl {
    public:
      Segment &s;
      int repeats;
      NeuronReal ratio;
      std::vector<Quality> qualities;
      
      H(Segment &s, int repeats, NeuronReal ratio): s(s), repeats(repeats), ratio(ratio) { }
      
      void threadFunc(Barrier &barrier) override {
        Segment &s = this->s;
        Layout l = s.layout;
        NeuronReal r = ratio;
        int sx = l.getW() - s.sx + 1;
        int sy = l.getH() - s.sy + 1;
        int sz = l.getD() - s.sz + 1;
        
        Quality q;
        for(int i = 0; i < repeats; ++i) {
          int x = l.x0 + barrier.commonRand()%sx;
          int y = l.y0 + barrier.commonRand()%sy;
          int z = l.z0 + barrier.commonRand()%sz;
          q += s.pass(barrier, x, y, z, r);
          barrier.wait();
        }
        qualities[barrier.tid] = q;
      }
    } h(s, repeats, backpass ? 1 : 0);
    
    int cnt = s.layout.getCount();
    values.resize(cnt);
    weights.resize(s.weightsCount);
    s.f_values = values.data();
    s.weights = weights.data();

    for(int i = 0; i < cnt; ++i)
      s.f_values[i] = rand()/(NeuronReal)RAND_MAX;
    for(int i = 0; i < s.weightsCount; ++i)
      s.weights[i].w = rand()/(WeightReal)RAND_MAX;
    h.qualities.resize(threads);
    
    volatile long long t0 = timeUs();
    h.runThreads(threads);
    volatile long long t1 = timeUs();

    AccumReal sum = 0;
    for(int i = 0; i < threads; ++i) sum += h.qualities[i].train + h.qualities[i].human;
    for(int i = 0; i < s.weightsCount; ++i) sum += s.weights[i].w;
    
    printf("%s %d: %f, %lld\n", name, (int)backpass, (t1 - t0)*1e-6, (long long)sum);
  }


  static void runSegment(const char *name, Segment &s, int threads, int repeats) {
    runSegment(name, s, threads, repeats, false);
    runSegment(name, s, threads, repeats, true);
  }

  
  static void runSegment(const char *name, Segment &s, int threads) {
    int repeats = (int)( (1ll*1000*1000*1000 + s.effectiveLinks/2)/s.effectiveLinks );
    if (!SegmentTest::testSegment(name, s)) return;
    runSegment(name, s, threads, repeats, false);
    runSegment(name, s, threads, repeats, true);
  }
  
  
  static void run(int threads) {
    { SegmentCx4 s(3, 24); s.layout = Layout(514, 514,  3).expandXY(2);
      runSegment("cx4-3x24", s, threads); }
  }
};


std::vector<NeuronReal> BenchmarkSegment::values;
std::vector<Weight> BenchmarkSegment::weights;