#include <iostream>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <sys/shm.h>
#include <sys/stat.h>
#include <sys/wait.h>

#include <sys/mman.h>
#include <unistd.h>

#include <vector>
#include <iomanip>
#include <cmath>

using namespace std;

enum class StateT { WORKER_READY, EXIT, WORKER_DATA, WORKER_DONE};
ostream & operator << (ostream & s, StateT state) {
    s << static_cast<int>(state) ;
    return s;
}

struct RecordT {
     int start,
         stop;
     StateT state;
     long double result;
};

int WORKERS{8};
int SIZE{static_cast<int>(sizeof(RecordT)) * WORKERS};

const int WORK_SIZE{10'000};
const int WORK_LIMIT{2'000'000'000};

// look at /dev/shm
const string NAME{"/PISharedMem"};

RecordT * MakeSharedMemory();
void InitStart(RecordT data[], int & limit) ;
void  MakeWorkers(RecordT data[]);

long double DoPI(int start, int end);
void Worker(RecordT array[], int myId);
void FinishWork(RecordT data[], int limit, long double & pi);
void FinalSum(RecordT data[], long double & pi);
void CleanUp(void);

int main( int argc , char ** argv ) {

    if (argc == 2) {
       int tmp = atoi(argv[1]);
       if (tmp >= 0) {
          WORKERS = tmp;
          SIZE = sizeof(RecordT) * WORKERS;
       }
    }

    cout << "Running with " << WORKERS << " workers." << endl << endl;

    long double pi = 4;

    if (WORKERS == 0) {
        pi += DoPI(0,WORK_LIMIT);   
    } else {

        RecordT * data {MakeSharedMemory()};

        if (data == nullptr) {
           cerr << "Exiting, could not set up shared memory" << endl;
           return 1;
        }

        int limit{0};
        InitStart(data, limit);

        MakeWorkers(data);

        // workers should not make it here.
        FinishWork(data, limit, pi);
        FinalSum(data, pi);
    }

    cout << setprecision(15) <<  "   PI is " <<  pi << endl;
    cout << setprecision(15) <<  " M_PI is " << M_PI << endl;

    CleanUp();

    return 0;
}

RecordT * MakeSharedMemory() {
    RecordT * data{nullptr};

    int openFlags {O_CREAT | O_RDWR | O_TRUNC};
    mode_t mode{ S_IRUSR | S_IWUSR};
    int fd{shm_open(NAME.c_str(), openFlags, mode)};

    if (fd == -1) {
       perror("\tError shm_open: ");
       return data;
    }

    if( -1 == ftruncate(fd, SIZE)) {
       perror("\tError ftruncate: ");
       shm_unlink(NAME.c_str());
       return data;
    }

    int protect{ PROT_READ| PROT_WRITE};
    data = reinterpret_cast<RecordT *> (
              mmap(0, SIZE, protect, MAP_SHARED, fd, 0));

    close(fd);

    if (data == MAP_FAILED) {
       perror("\tError mmap");
       data = nullptr;
    } 

    return data;
}

void  MakeWorkers(RecordT data[]){
    int myID {-1};

    for(int proc = 0; proc < WORKERS; ++proc) {
       pid_t forkID{fork()};
       switch (forkID) {
          case -1 :
              perror("\tError, fork: ");
              exit(1);
          case 0:
              myID = proc;
              Worker(data, myID);
              // worker will call exit, but just in case.
              exit(0);
          default:
               break;
       }
    }
}

// wait for all the workers to exit then clean up the memory
void CleanUp(){
    int status;

    // wait for all of the workers
    for(int i = 0; i < WORKERS; ++i) {
       wait(&status);
    }

    // free the shared memory segment.
    if( -1 == shm_unlink(NAME.c_str())) {
        perror("\tError shm_unlink");
    }
}

// set up the inital work before we create the workers
void InitStart(RecordT data[], int & limit)  {
    for(int i = 0; i < WORKERS; ++i) {
        data[i].start = i * WORK_SIZE;
        data[i].stop = (i+1) * WORK_SIZE - 1;
        data[i].state = StateT::WORKER_DATA;
        limit = (i+1) * WORK_SIZE;
    }
}

// send work to workers that have finished the inital computation
void FinishWork(RecordT data[], int limit, long double & pi){

    while (limit < WORK_LIMIT) {
       for(int i = 0; i < WORKERS; ++i) {
          if (data[i].state == StateT::WORKER_DONE and limit <= WORK_LIMIT) {
              // get the result
              pi += data[i].result;

              // write the next case
              data[i].start = limit;
              data[i].stop = limit + WORK_SIZE -1;

              // do this last, it is my form of synch control
              data[i].state = StateT::WORKER_DATA;

              limit += WORK_SIZE;
          }
       }
    }

}

// collect the final values 
void FinalSum(RecordT data[], long double & pi) {

    int finished{0};

    while (finished < WORKERS) {
       for(int i = 0; i < WORKERS; ++i) {
          if (data[i].state == StateT::WORKER_DONE) {
              pi += data[i].result;
              ++finished;
              data[i].state = StateT::EXIT;
          }
       }
    }
    return;
}

// worker compute pi
long double DoPI(int start, int end) {
    int sign{1};

    if (start %2 ==  0) {
      sign = -1; 
    }

    long double sum{0};
    long double term{(start+1)*2.0 + 1.0};

    for(int i = start; i <= end; ++i){
        sum = sum + sign * 4.0/term;
        sign = - sign;
        term += 2;
    }

    return sum;
}

void Worker(RecordT data[], int myID) {
     bool done{false};
     bool first{true};

     usleep(10);
     while (not done) {
         if (first) {
            usleep(0);
         } else {
            first=false;
         }
        
         volatile StateT state = data[myID].state;
         switch (state) {
            case StateT::WORKER_DATA:
              /* cout << "Worker " << myID << " is working on "
                    << data[myID].start << " to " 
                    << data[myID].stop << endl;
*/
               // calculate the result
               data[myID].result = DoPI(data[myID].start, data[myID].stop);

               // do this last, this is my sync chontrol
               data[myID].state = StateT::WORKER_DONE;
               break;
            case StateT::EXIT:
               done = true;
               break;
            default:
               break;
         }
     }
     exit(0);
}