#include <stdio.h>
#include <GL/glut.h>
#include <mpi.h>

#define SIZE 4 //square of the max distance from origin
#define A0 0.0
#define B0 0.0
#define X 800
#define Y 800
#define ITER 10000
int plot[X][Y];

long double x_width=4.0;
long double y_width=4.0;
long double start_x=-2.0;
long double start_y=-2.0;

int left_x, right_x, top_y, bottom_y;

int iterate(long double x, long double y) {
    int i, j;
    long double temp, a, b;

    a=A0;
    b=B0;
    
    for(j=0;j<ITER;j++) {
        if(a*a+b*b > SIZE) break;
        
        temp=2.0*a*b;
        a*=a;
        a-=b*b;
        b=0.0;
        b=temp+y;
        a+=x;
    }
    if(j==ITER) j=0;
    return j;
}

void display(void) {
    int i,j;
    glClear(GL_COLOR_BUFFER_BIT);
    glBegin(GL_POINTS);
    for(i=0;i<X;i++) {
        for(j=0;j<Y;j++) {
            glColor3f(plot[i][j]%13/16.0, plot[i][j]%47/16.0, plot[i][j]%31/16.0);
            glVertex3f(i*1.0, j*1.0, 0.0);
        }
    }

    glEnd();
    glFlush();
    glutSwapBuffers();
}


void myinit(void) {
   glClearColor(0.0, 0.0, 0.0, 0.0);

   glMatrixMode(GL_PROJECTION);
   glLoadIdentity();
   glOrtho(0.0, X*1.0, 0.0, Y*1.0, -1.0, 1.0);
}

int abs(int n) {
    return (n>=0 ? n : -1*n);
}

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

int next_work=0, num_nodes;
int my_rank;
int signed_out[256];
int i, j, assn;
int line_buf[Y+1];
int in_count=0;
MPI_Status stat;

MPI_Init(&argc, &argv);
MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

if(my_rank == 0) { 

    glutInit(&argc, argv);
    glutInitDisplayMode(GLUT_DOUBLE|GLUT_RGB);
    glutInitWindowSize(X, Y);
    glutInitWindowPosition(300,50);
    glutCreateWindow("Mandelbrot w/MPI");
    myinit();

    glutDisplayFunc(display);

    MPI_Comm_size(MPI_COMM_WORLD, &num_nodes);
    
    for(i=1;i<num_nodes;i++) { //loop thru the nodes (not the master)
        MPI_Send(&next_work, 1, MPI_INT, i, 0, MPI_COMM_WORLD);
        //printf("MASTER: sent row %d to node %d\n", next_work, i);
        signed_out[i]=next_work;
        next_work++;
    }
    
   while(in_count < Y) {
        MPI_Recv(&line_buf, Y+1, MPI_INT, MPI_ANY_SOURCE, 0, MPI_COMM_WORLD, &stat);
        //printf("Got work in from %d\n", line_buf[Y]);
        in_count++;
        for(j=0;j<Y;j++) {
            plot[signed_out[line_buf[Y]]][j] = line_buf[j];
        }
        glutPostRedisplay();
if(next_work < Y) {
        MPI_Send(&next_work, 1, MPI_INT, line_buf[Y], 0, MPI_COMM_WORLD);
         //printf("MASTER: sent row %d to node %d\n", next_work, line_buf[Y]);
        signed_out[line_buf[Y]]=next_work;
        next_work++;    
     }
   } 
    next_work=-1;
    
    for(i=1;i<num_nodes;i++) { //loop thru the nodes (not the master)
        MPI_Send(&next_work, 1, MPI_INT, i, 0, MPI_COMM_WORLD);
    }
    
    glutPostRedisplay();
    glutMainLoop();
}

else { //slave process
    assn=-2;
    while(assn!=-1) {
        MPI_Recv(&assn, 1, MPI_INT, 0, 0, MPI_COMM_WORLD, &stat);
        //printf("Node %d: Received Work\n", my_rank);
        for(j=0;j<Y;j++) {
            line_buf[j] = iterate(x_width*(1.0*assn/X)+start_x, y_width*(1.0*j/Y)+start_y);
        }
        line_buf[j] = my_rank;
        //printf("1\n");
        MPI_Send(&line_buf, Y+1, MPI_INT, 0, 0, MPI_COMM_WORLD);    
        //printf("2\n");
    }
}

MPI_Finalize();

}