#include <math.h>
#include <mpi.h>
#include <stdbool.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define XSIZE 2560
#define YSIZE 2048

#define MAXITER 255

double xleft = -2.01;
double xright = 1;
double yupper, ylower;
double ycenter = 1e-6;
double step;

int pixel[XSIZE * YSIZE];

#define PIXEL(i, j) ((i) + (j) * XSIZE)

typedef struct {
  double real, imag;
} complex_t;

int rank, size;
bool save = false;

void calculate() {
  for (int i = 0; i < XSIZE; i++) {
    for (int j = 0; j < YSIZE; j++) {
      /* Calculate the number of iterations until divergence for each pixel.
         If divergence never happens, return MAXITER */
      complex_t c, z, temp;
      int iter = 0;
      c.real = (xleft + step * i);
      c.imag = (ylower + step * j);
      z = c;
      while (z.real * z.real + z.imag * z.imag < 4) {
        temp.real = z.real * z.real - z.imag * z.imag + c.real;
        temp.imag = 2 * z.real * z.imag + c.imag;
        z = temp;
        if (++iter == MAXITER)
          break;
      }
      pixel[PIXEL(i, j)] = iter;
    }
  }
}

typedef unsigned char uchar;

/* save 24-bits bmp file, buffer must be in bmp format: upside-down */
void savebmp(char *name, uchar *buffer, int x, int y) {
  FILE *f = fopen(name, "wb");
  if (!f) {
    printf("Error writing image to disk.\n");
    return;
  }
  unsigned int size = x * y * 3 + 54;
  uchar header[54] = { 'B', 'M', size & 255, (size >> 8) & 255, (size >> 16) & 255, size >> 24, 0,
                       0, 0, 0, 54, 0, 0, 0, 40, 0, 0, 0, x & 255, x >> 8, 0, 0, y & 255, y >> 8, 0, 0, 1, 0, 24, 0, 0, 0, 0, 0, 0,
                       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };
  fwrite(header, 1, 54, f);
  fwrite(buffer, 1, XSIZE * YSIZE * 3, f);
  fclose(f);
}

/* given iteration number, set a colour */
void fancycolour(uchar *p, int iter) {
  if (iter == MAXITER)
    ;
  else if (iter < 8) {
    p[0] = 128 + iter * 16;
    p[1] = p[2] = 0;
  } else if (iter < 24) {
    p[0] = 255;
    p[1] = p[2] = (iter - 8) * 16;
  } else if (iter < 160) {
    p[0] = p[1] = 255 - (iter - 24) * 2;
    p[2] = 255;
  } else {
    p[0] = p[1] = (iter - 160) * 2;
    p[2] = 255 - (iter - 160) * 2;
  }
}

int main(int argc, char **argv) {
  if (argc <= 1) {
    printf("usage: mandel_mpi [y/n]\n");
    printf("y/n denotes if you want to save the resulting\n");
    printf("calculations to mandel.bmp or not.\n");
    return 1;
  } else {
    char arg = argv[1][0];
    save = arg == 'y' || arg == '1';
  }
  MPI_Init(&argc, &argv);
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  MPI_Comm_size(MPI_COMM_WORLD, &size);

#undef YSIZE
#define YSIZE (2048 / size)

  step = (xright - xleft) / XSIZE;
  double old_yupper = ycenter + (step * YSIZE * size) / 2;
  double old_ylower = ycenter - (step * YSIZE * size) / 2;
  double delta = (old_yupper - old_ylower) / size;

  ylower = old_ylower + rank * delta;
  yupper = old_ylower + (rank + 1) * delta;
  ycenter = (ylower + yupper) / 2;

  calculate();

  printf("after calculate %d\n", rank);

  const int bufsize = XSIZE * YSIZE * 3;
  unsigned char *buffer = calloc(bufsize, 1);
  for (int i = 0; i < XSIZE; i++) {
    for (int j = 0; j < YSIZE; j++) {
      int p = ((YSIZE - j - 1) * XSIZE + i) * 3;
      fancycolour(buffer + p, pixel[PIXEL(i, j)]);
    }
  }

  printf("after buffer %d\n", rank);

  if (rank != 0)
    MPI_Send(buffer, bufsize, MPI_UNSIGNED_CHAR, 0, 0, MPI_COMM_WORLD);

  printf("hello again from %d\n", rank);

  // use process 0 as central process: gather calulations and output image
  if (rank == 0) {
    unsigned char *total = calloc(bufsize, size);
    memcpy(total + (size - 1) * bufsize, buffer, bufsize);
    printf("after memcpy 0\n");

    for (int i = 1; i < size; i++) {
      MPI_Recv(total + (size - 1 - i) * bufsize, bufsize, // write reverse order
               MPI_UNSIGNED_CHAR, i, 0, MPI_COMM_WORLD, MPI_STATUS_IGNORE);
      printf("after direct receive %d\n", i);
    }
    printf("after loop\n");

#undef YSIZE
#define YSIZE 2048

    if (save)
      savebmp("mandel.bmp", total, XSIZE, YSIZE);
    printf("after save %d\n", save);
  }

  MPI_Finalize();
  return 0;
}