#include <iostream>
#include <cstdlib>
 
#include <ff/pipeline.hpp>
#include <ff/farm.hpp>
 
#define DELAYSUM 100000
#define DELAYSQ  500000
 
#define MAX(a,b) (a>b ? a : b)
 
using namespace std; 
using namespace ff; 
 
class Iota : public ff_node {
private: 
  int m; 
public: 
  Iota(int m):m(m) {}
 
  void * svc(void * t) {
    for(int i=1; i<=m; i++) {
      int * n = new int(i); 
      ff_send_out((void *) n); 
    }
    return((void *) FF_EOS);
  }
};
 
class Summer : public ff_node {
private: 
  int sum;
public: 
  int svc_init() {
    sum = 0;
    return(0); 
  }
 
  void * svc(void * t) {
    sum += *((int *) t);
    return (GO_ON); 
  }
 
  void svc_end() {
    cout << "sum = " << sum << endl; 
    return;
  }
};
 
class Incr : public ff_node {
  void * svc(void * t) {
    int * r = new int( ++( *((int *) t)) );
    free(t);
    usleep(DELAYSUM);
    return ((void *) r); 
  }
}; 
class Square : public ff_node {
  void * svc(void * t) {
    int x = *((int *) t); 
    free(t);
    usleep(DELAYSQ);
    return ((void *) new int(x*x)); 
  }
}; 
 
// a.out maxnum nw
int main(int argc, char * argv[]) {
  int m = atoi(argv[1]); 
  int nw= atoi(argv[2]); 
 
  std::cout << "Computing   " << m << " tasks" << std::endl;;
  std::cout << "DELAYSUM is " << DELAYSUM/1000 << " msecs" << endl;;
  std::cout << "DELAYSQ  is " << DELAYSQ/1000 << " msecs" << endl;;
 
  // build main pipeline : three stages(generator, farm(pipe(inc,sq)), summer)
 
  ff_pipeline mainPipe; 
  mainPipe.add_stage(new Iota(m)); // generator
 
  // second stage: farm(pipe(inc, seq))
  // as much workers as the parameter from the command line
  ff_farm<> farmstage; 
  vector<ff_node*> w; 
  for(int i=0; i<atoi(argv[2]); i++) {
    ff_pipeline * p = new ff_pipeline(); 
    p->add_stage(new Incr); 
    p->add_stage(new Square);
    w.push_back(p); 
  }
  farmstage.add_workers(w);
  farmstage.add_collector(NULL);   // adds default collector
                                   // default emitter already here 
  mainPipe.add_stage(&farmstage); 
 
  mainPipe.add_stage(new Summer); 
 
  mainPipe.run_and_wait_end(); 
 
  std::cout << "Spent " << mainPipe.ffTime() << " msecs " << std::endl;
  return(0); 
}