#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/ipc.h>
#include <sys/msg.h>
#include <sys/wait.h>
#include <string.h>

#define KEY 1234
#define NO_MESSAGES_BYTES 8000
#define DEFAULT_WORKERS 4

/*
 * Para os cálculos funcionarem, são enviadas um total de
 * 1000 mensagens (referente a 4 workers)
 * com 8000 bytes cada, dos filhos para o pai
 */

struct tuplo {
		int pid;
		int  nx, ny;
		double xmin, xmax, ymin, ymax;
};

typedef struct {
	long mtype;
	struct tuplo mtext;
} mensagem;

typedef struct {
	long mtype;
	int part; // indica em cada mensagem qual é a parte de um total de mensagens que o worker envia
	double buffer[1000*1000*sizeof(double)/NO_MESSAGES_BYTES];
} calculos;

int max_iterations = 256;

/* This function (compute_point) is the algorithmic core of this program,
   computing what could be called the Mandelbrot function of the complex
   number (cr, ci).
*/

double compute_point(double ci, double cr) {
  int iterations = 0;
  double zi = 0;
  double zr = 0;
  
  while ((zr*zr + zi*zi < 4) && (iterations < max_iterations)) {
    double nr, ni; 

    /* Z <-- Z^2 + C */
   
    nr = zr*zr - zi*zi + cr;
    ni = 2*zr*zi + ci;
    
    zi = ni;
    zr = nr;

    iterations ++;
  }  
  return iterations;  
}

/* The "compute" function computes the Mandelbrot function over every
   point on a grid that is "nx" points wide by "ny" points tall, where
   (xmin,ymin) and (xmax,ymax) give two corners of the region the 
   complex plane.
*/

void compute(double *buffer, int nx, int ny, double xmin, double xmax, 
	     double ymin, double ymax) {
  double delta_x, delta_y;
  int x, y;

  delta_x = (xmax - xmin)/nx;
  delta_y = (ymax - ymin)/ny;

  for (y=0;  y<ny; y++) {
    double y_value = ymin + delta_y * y;
    for (x=0; x<nx; x++) {
      double x_value = xmin + delta_x * x;
      buffer[y*nx + x] = compute_point(x_value, y_value);
    }
  }
}

/* Output the data contained in the buffer to a Portable Greymap format
   image file.  The parameter "max" should be an upper bound for the
   data values in the buffer. 
*/

void output_pgm(char *filename,double *buffer, int nx, int ny, double max) {
  int i;
  FILE *file;
  file = fopen(filename,"w");
  fprintf(file,"P2\n");
  fprintf(file,"%d %d\n",nx,ny);
  fprintf(file,"%d\n",(int)max);
  for (i=0; i<nx*ny; i++) {
    if (!(i%nx)) fprintf(file,"\n");
    fprintf(file,"%d ",(int)buffer[i]);
  }
  fclose(file);
}

void tcompute(int queue) {
	mensagem tmp;
	if(msgrcv(queue, &tmp, sizeof(struct tuplo), getpid(), 0) == -1) perror("msgrcv");

	calculos snd;
	snd.mtype = getpid() * 2; // para cada worker ter um número diferente ao enviar a resposta
	snd.part = 0;

	double *buf;
	buf = malloc(tmp.mtext.nx * tmp.mtext.ny * sizeof(double)); // cria-se um buffer temporario
	compute(buf, tmp.mtext.nx, tmp.mtext.ny, tmp.mtext.xmin, tmp.mtext.xmax, tmp.mtext.ymin, tmp.mtext.ymax);

	int no_messages = (tmp.mtext.nx * tmp.mtext.ny * sizeof(double)) / NO_MESSAGES_BYTES; // números de mensagens que o worker tem que enviar
	int i;
	for(i = 0; i < no_messages; i++) { // para cada mensagem
		memcpy(snd.buffer, buf, NO_MESSAGES_BYTES); // copia-se os bytes do buffer temporario para a estrutura
		if(msgsnd(queue, &snd, sizeof(int) + NO_MESSAGES_BYTES, 0) == -1) perror("msgsnd");
		snd.part++;
		buf = buf + 1000; // anda-se 1000 bytes para a frente no buffer temporario
	}
}

int main(int argc, char *argv[]){
	int workers = (argc > 1) ? atoi(argv[1]) : DEFAULT_WORKERS;
	pid_t pids[workers];
	int queue = msgget(KEY, IPC_CREAT|0600); // cria-se a message queue

	double *b, *ptr;
	b = ptr = malloc(1000*1000*sizeof(double));

	int t;
	pid_t pid;
	for(t = 0; t < workers; t++) { // lança os workers
		pid = fork();
		if(pid == 0) {
			tcompute(queue); exit(0);
		}
		pids[t] = pid;
	}

	int part = 1000 / workers;
	double img_part = 2.0 / workers;
	int i; 
	for(i = 0; i < workers; i++) { // faz os calculos para saber o que os workers têm que fazer e envia as mensagens
		mensagem tmp;

		tmp.mtype = pids[i];
		tmp.mtext.pid = pids[i];
		tmp.mtext.nx = 1000;
		tmp.mtext.ny = part;
		tmp.mtext.xmin = -1.0;
		tmp.mtext.xmax = 1.0;
		tmp.mtext.ymin = -1.0 + img_part * i;
		tmp.mtext.ymax = -1.0 + img_part + img_part * i;

		if(msgsnd(queue, &tmp, sizeof(struct tuplo), IPC_NOWAIT) == -1) perror("msgsnd");
	}

	int no_messages = workers * ((1000 * part * sizeof(double)) / NO_MESSAGES_BYTES); // número de mensagens total que o master recebe
	int received = 0;

	calculos tmp;

	while(received < no_messages) {
		/* enquanto não receber todas as mensagens tenta obter mensagens de cada worker em particular
		 * mas não bloqueia à espera de resposta, em vez disso se não recebe de um, tenta de outro
		 */
		for(i = 0; i < workers; i++) {
			if(msgrcv(queue, &tmp, sizeof(int) + NO_MESSAGES_BYTES, pids[i] * 2, IPC_NOWAIT) != -1) { // se recebeu
				b = ptr + (i * part * 1000) + tmp.part * 1000; // mete o b a apontar para o sitio onde a parte que se recebeu vai ser copiada
				memcpy(b, tmp.buffer, NO_MESSAGES_BYTES); // copia a parte de dados recebida para o buffer da imagem
				received++;
			}
		}
	}

	/* espera por todos os workers */
	for(i = 0; i < workers; i++)
		wait(NULL);

	/* fecha a message queue e faz o output da imagem */
	msgctl(queue, IPC_RMID, NULL);
	output_pgm("mandel.pgm", ptr, 1000, 1000, 255);
	return 0;
}