#include <iostream>

#include <iomanip>

#include <boost/mpi.hpp>
#include <boost/mpi/environment.hpp>
#include <boost/mpi/communicator.hpp>

#include <queue>
#include <vector>
#include <optional>

namespace mpi = boost::mpi;

using namespace std;

// could be an enum, but MPI would probably not like that 
const int WORK_REQUEST_TAG{1};
const int WORK_DONE_TAG{2};
const int WORK_TAG{3};
const int NO_MORE_WORK_TAG{4};
const int WORK_DATA_TAG{5};
const int WORK_AVAILABLE_TAG{6};

// Some useful aliases.
using BlockT  = vector<int>;
using ImageT = vector<vector<int>>;
using CommT = mpi::communicator;
using ReqListT = vector<mpi::request>;

// image information
// height and width will be mutiples of the block size for now.
size_t pictureWidth{40};
size_t pictureHeight{80};
size_t workBlockX{10};
size_t workBlockY{10};

// plane information
//double ly{-1},lx{-2}, uy{1}, ux{1};
//double lx{-1.1},ly{-0.1}, ux{0.1}, uy{1};
double lx{-1.5},ly{-0.1}, ux{-1}, uy{0.5};
double dx = (ux - lx)/static_cast<float>(pictureWidth); 
double dy = (uy - ly)/static_cast<float>(pictureHeight); 

// struct for data transfer
struct GridPoint {
    public:
        friend class boost::serialization::access;

        template<typename Archive> 
        void serialize(Archive & ar, const unsigned int version) { 
            ar & x;
            ar & y;
            ar & row;
            ar & col;
        };

        double x, 
               y;
        size_t row, 
               col;
};

using WorkQT = queue<GridPoint>;
using WorkListT = vector<GridPoint>;

void MakeWorkQueue(WorkQT & work);
void DoInitialListen(CommT world, ReqListT & reqs);
void ProcessMessage(CommT world, ReqListT & reqs, int & activeJobs, WorkQT & work, WorkListT & assignedWork, ImageT & image);
void SendFinalNotice(CommT world, ReqListT & reqs);

void ProcessResult(BlockT data, GridPoint p, ImageT & image);
void PrintImage(const ImageT image);

void WorkerMain(CommT world);
BlockT ProcessWork(CommT world, const GridPoint & p);

int main([[maybe_unused]] int argc, [[maybe_unused]] char** argv) {

    mpi::environment env(argc, argv);
    CommT world;

    if (world.rank() == 0) { 
        ImageT image(pictureWidth, vector<int>(pictureHeight,0)); 
        WorkListT assignedWork(world.size());
        WorkQT work;
     
        int activeJobs{0};
        ReqListT reqs(world.size());

        MakeWorkQueue(work);

        DoInitialListen(world, reqs);

        while (work.size() > 0 or activeJobs > 0) {
           ProcessMessage(world, reqs, activeJobs, work, assignedWork, image);
         }

        SendFinalNotice(world, reqs);
        PrintImage(image);

    } else {
        WorkerMain(world);
        return 0;
    }
    return 0;
}

void PrintImage(const ImageT image) {
    cout << "THE IMAGE " << endl;
    cout << "it is " << image[0].size() << " x " << image.size() << endl;
    for(auto & row: image) {
       for (auto col: row) {
          if (col > 0) {
             cout << "*";
          } else {
             cout << " ";
          }
       }
       cout << endl;
    }
    cout << endl;
}

void MakeWorkQueue(WorkQT & work) {
    size_t row,col;

    for(col = 0; col < pictureWidth; col += workBlockX) {
       for(row = 0; row < pictureHeight; row +=  workBlockY) {
          GridPoint p{lx+dx*col, ly+dy*row,row,col};
          work.push(p);
       }
    }
}

void DoInitialListen(CommT world, ReqListT & reqs) {
    int i = 1;
    // set a listener for each workier
    for(i = 1; i < world.size(); ++i) {
        reqs[i] = world.irecv(i, MPI_ANY_TAG);
    }
}

void ProcessMessage(CommT world, ReqListT & reqs, int & activeJobs,
                    WorkQT & work, WorkListT & assignedWork, ImageT & image){
    boost::optional<mpi::status> status;
    BlockT result;
    int i;

    for(i = 1; i < world.size(); ++i) {
        if (reqs[i].active() and  (status = reqs[i].test()) ) {
           int tag = status.value().tag();
           int source = status.value().source();

           switch(tag) {
               case  WORK_DONE_TAG:
                  world.recv(source, WORK_DATA_TAG, result);

                  ProcessResult(result, assignedWork[source], image);
                  --activeJobs;
                  reqs[i] = world.irecv(i, MPI_ANY_TAG);

                  break;

               case WORK_REQUEST_TAG:
                  if (work.size() > 0) {
                     world.send(source, WORK_AVAILABLE_TAG);

                     // send the work
                     GridPoint p{work.front()};
                     work.pop();
                     world.send(source, WORK_TAG, p);
                     assignedWork[source] = p;
                     ++activeJobs;

                     // we will get another request from this client
                     reqs[i] = world.irecv(i, MPI_ANY_TAG);
                  } else {
                     world.send(source, NO_MORE_WORK_TAG);
                  }
                  break;
           }
        }
    }
}

void SendFinalNotice(CommT world, ReqListT & reqs) {
    int i;

    for(i = 1; i < world.size(); ++i) {
       if (reqs[i].active()) {
          reqs[i].wait();
          world.send(i, NO_MORE_WORK_TAG);
       }
    }
}

void ProcessResult(BlockT data, GridPoint p, ImageT & image){
    for (size_t col = 0; col < workBlockX; ++col) {
        for(size_t row = 0; row < workBlockY; ++row) {
            image[col+p.col][row+p.row] = data[row*workBlockX + col];
        }
    }
}

// worker routines.
//
void WorkerMain(CommT world) {
     GridPoint p;
     BlockT b;
     mpi::status status;

     // initial work request
     world.send(0, WORK_REQUEST_TAG);
     status = world.recv(0, MPI_ANY_TAG);

     while (status.tag() == WORK_AVAILABLE_TAG) {
         world.recv(0, WORK_TAG, p);

         b = ProcessWork(world, p);

         world.send(0, WORK_DONE_TAG);
         world.send(0, WORK_DATA_TAG, b);

         // request more work
         world.send(0, WORK_REQUEST_TAG);
         status = world.recv(0, MPI_ANY_TAG);
     }
}

const int MAX_ITERATIONS{10000};

int Iterate(long double x, long double y) {
    complex<long double> z{x,y};
    complex<long double> c{x,y};

    int iterations{0};
       
    while (norm(z) < 2 and iterations < MAX_ITERATIONS) {
       z = z * z + c;
       ++iterations;
    }

    if (iterations == MAX_ITERATIONS) {
       return 1;
    } else {
       return 0;
    }

    return iterations;

}

BlockT ProcessWork(CommT world, const GridPoint & p) {
        BlockT data(workBlockX*workBlockY, world.rank());
        
        size_t row,col;
        long double x,y;
        int iters;

        y = p.y;
        for(row = 0; row < workBlockY; ++row) {
           x = p.x;
           for( col = 0; col < workBlockX; ++col) {
              iters = Iterate(x,y);     
              x += dx;
              data[row*workBlockX + col] = iters;
           }
           y += dy;
        }
       
        return data;
}