//
// Starter MPI Mandelbrot program.
//
// This program computes and displays all or part of the Mandelbrot 
// set.  By default, it examines all points in the complex plane
// that have both real and imaginary parts between -2 and 2.  
// Command-line parameters allow zooming in on a specific part of
// this range.
// 
// Usage:
//   mandelbrot [x y size] [maxiter]
// where 
//   x, y, and size specify the range to examine (a square of size
//     "size", centered at x + iy -- by default, a square of size 4
//     centered at the origin)
//   maxiter denotes the maximum number of iterations at each
//     point (default MAX_ITER, see below)
// Arguments are optional.
// 
// Input:  none, except the optional command-line arguments
// Output: a graphical display as described in Wilkinson & Allen,
//   displayed using the X Window system, plus text output to
//   standard output showing the above parameters, plus execution
//   time in seconds.
// 
// 
// Original code obtained from Web site for Wilkinson and Allen's
// text on parallel programming:
// http://www.cs.uncc.edu/~abw/parallel/par_prog/
// 
// Reformatted and revised by B. Massingill, Jan 2001 and Feb 2002
// to produce revised sequential version.
//
// Additional changes made for MPI.
//

#include <iostream>
#include <cstring>
#include <cmath>
#include <cstdlib>
#include <mpi++.h>
#include <unistd.h>
#include <X11/Xlib.h>
#include <X11/Xutil.h>
#include <X11/Xos.h>

// Structure definition for complex numbers
// (We could use a class, but this is noticeably faster.)
typedef struct {
  double real, imag;
} complex ;

// Shorthand for some commonly-used types
typedef unsigned int uint;
typedef unsigned long ulong;

// ---- Function declarations ------------------------------------------

int master_pgm(const int nslaves, const uint width, const uint height,
	       const double real_min, const double real_max,
	       const double imag_min, const double imag_max,
	       const uint maxiter);
int slave_pgm(const int myID, const uint width, const uint height,
	      const double real_min, const double real_max,
	      const double imag_min, const double imag_max,
	      const uint maxiter);
int setup(const uint width, const uint height,
	  Display * &display, Window &win, GC &gc, 
	  ulong &min_color, unsigned long &max_color);
void interact(Display * &display, Window &win, 
	      const uint width, const uint height,
	      const double real_min, const double real_max,
	      const double imag_min, const double imag_max);

// ---- Main program ---------------------------------------------------

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

// Initialize for MPI

  int nprocs;
  int myid;
  int returnval;

  if (MPI_Init(&argc, &argv) != MPI_SUCCESS) {
    cerr << "MPI initialization error\n"; 
    exit(EXIT_FAILURE);
  }

  MPI_Comm_size(MPI_COMM_WORLD, &nprocs);
  MPI_Comm_rank(MPI_COMM_WORLD, &myid);

  if (nprocs < 2) {
    cerr << "Number of processes must be at least 2\n";
    MPI_Finalize();
    exit(EXIT_FAILURE);
  }

  const uint width = 800;	// dimensions of display window,
  const uint height = 800;	//   in pixels

  // Check command-line arguments, if any
  // Usage is:
  //   mandelbrot [x y size] [maxiter]
  //   or
  //   mandelbrot -h to show usage message

  if (((argc > 1) && (strcmp(argv[1], "-h") == 0))
      || ((argc > 1) && (argc < 4))) {
    if (myid == 0)
      cerr << "Usage:  mandelbrot [x y size] [maxiter]\n";
    MPI_Finalize();
    exit(EXIT_FAILURE);
  }

  // Set default values
  uint maxiter = 100;	// maximum iterations
  double			// range to examine
    real_min = -2, real_max = 2,
    imag_min = -2, imag_max = 2;

  // Process command-line arguments, if any
  if (argc > 1) {
    // argc >= 4 -- we checked this earlier
    double x = atof(argv[1]);
    double y = atof(argv[2]);
    double size = atof(argv[3]);
    real_min = x - size/2;
    real_max = real_min + size;
    imag_min = y - size/2;
    imag_max = imag_min + size;
  }
  if (argc > 4) 
    maxiter = atoi(argv[4]);
 
  // Call master or slave code as appropriate

  if (myid == 0) {
#ifdef DEBUG
    cerr << "master process started\n";
#endif // DEBUG
    returnval = 
      master_pgm(nprocs-1, width, height,
		 real_min, real_max, imag_min, imag_max, maxiter);
  }
  else {
#ifdef DEBUG
    cerr << "slave process started, id = " << myid << endl;
#endif // DEBUG
    returnval = 
      slave_pgm(myid, width, height,
		real_min, real_max, imag_min, imag_max, maxiter);
  }

  // Finish up

  MPI_Finalize();

  return returnval;
}
// ---- Program for master process -------------------------------------

// Contains most code from original main program, slightly modified
//   for MPI 
// To be called in process 0 only

// pre:  all parameters have values as assigned and described in "main"
// post:  returns EXIT_SUCCESS or EXIT_FAILURE as appropriate
int master_pgm(const int nslaves, const uint width, const uint height,
	       const double real_min, const double real_max,
	       const double imag_min, const double imag_max,
	       const uint maxiter) {

  // Initialize for graphical display

  Display *display;
  Window win;
  GC gc;
  ulong min_color, max_color;

  if (setup(width, height, display, win, gc, min_color, max_color) !=
      EXIT_SUCCESS)
    MPI_Abort(MPI_COMM_WORLD, EXIT_FAILURE);
 
  // Start timing 
  double start_time = MPI_Wtime();

  // START OF SECTION TO CHANGE

  // Compute factors to scale computational region to window
  double scale_real = (double) (real_max - real_min) / (double) width;
  double scale_imag = (double) (imag_max - imag_min) / (double) height;

  // Compute factor for color scaling
  double scale_color = (double) (max_color - min_color) / 
		       (double) (maxiter - 1);

  // Calculate and draw points 

  for (uint j = 0; j < height; ++j) 
    for (uint i = 0; i < width; ++i) {

      double lengthsq, temp;
      complex z, c;
      z.real = z.imag = 0;

      // Scale display coordinates to actual region
      c.real = real_min + ((double) i * scale_real);
      c.imag = imag_min + ((double) (height-1-j) * scale_imag);
				// height-1-j so y axis displays
				// with larger values at top

      uint k = 0;
       // Calculate z0, z1, .... until divergence or maximum iterations
      do  {
	temp = z.real*z.real - z.imag*z.imag + c.real;
	z.imag = 2.0*z.real*z.imag + c.imag;
	z.real = temp;
	lengthsq = z.real*z.real + z.imag*z.imag;
	++k;
      } while (lengthsq < 4.0 && k < maxiter);

      // Scale color and display point 
      ulong color = (ulong) ((k-1) * scale_color) + min_color;
      XSetForeground (display, gc, color);
      XDrawPoint (display, win, gc, i, j);
    }

  // END OF SECTION TO CHANGE 

  // End timing 
  double end_time = MPI_Wtime();

  XFlush (display);

  // Produce text output 
  cout << "number of slave processes = " << nslaves << endl;
  cout << "center = (" << (real_max + real_min)/2
       << ", " << (imag_max + imag_min)/2
       << "), size = " << (real_max - real_min) << endl;
  cout << "maximum iterations = " << maxiter << endl;
  cout << "execution time in seconds = "
       << end_time - start_time << endl;

  // Wait for user response, then exit program 
  interact(display, win, width, height,
	   real_min, real_max, imag_min, imag_max);

  return EXIT_SUCCESS;
}

// ---- Program for slave process --------------------------------------

// To be called in each process other than process 0

// pre:  all parameters have values as assigned and described in "main"
// post:  returns EXIT_SUCCESS or EXIT_FAILURE as appropriate
int slave_pgm(const int myID, const uint width, const uint height,
	      const double real_min, const double real_max,
	      const double imag_min, const double imag_max,
	      const uint maxiter) {

  // START OF SECTION TO CHANGE
  // (put your code here) 
  // END OF SECTION TO CHANGE

  return EXIT_SUCCESS;
}

// ---- Other function definitions -------------------------------------

// initialize for graphical display:
// pre:  width, height are dimensions of display, in pixels
// post:  remaining parameters have been set to values needed by
//          caller
//        return value is EXIT_SUCCESS if all goes well,
//          EXIT_FAILURE otherwise 
int setup(const uint width, const unsigned int height,
	  Display * &display, Window &win, GC &gc, 
	  ulong &min_color, unsigned long &max_color) {

  // Variables for graphical display

  uint
    x = 0, y = 0,			// window position
    border_width = 4,			// border width in pixels
    disp_width, disp_height,		// size of screen
    screen;				// which screen

  char *window_name = "Mandelbrot Set",
    *disp_name = NULL;
  ulong valuemask = 0;
  XGCValues values;

  ulong white, black;	// white, black pixel values
 
  // Connect to Xserver

  if ( (display = XOpenDisplay (disp_name)) == NULL ) {
    cerr << "Cannot connect to X server "
	 << XDisplayName (disp_name) << endl;
    return EXIT_FAILURE;
  }
        
  /// Initialize for graphical display 

  screen = DefaultScreen (display);
  disp_width = DisplayWidth (display, screen);
  disp_height = DisplayHeight (display, screen);

  win = XCreateSimpleWindow (display, RootWindow (display, screen),
			     x, y, width, height, border_width, 
			     BlackPixel (display, screen), 
			     WhitePixel (display, screen));

  XStoreName(display, win, window_name);

  gc = XCreateGC (display, win, valuemask, &values);
				// graphics context
  white = WhitePixel (display, screen);	// color value for white
  black = BlackPixel (display, screen);	// color value for black

  XSetBackground (display, gc, white);
  XSetForeground (display, gc, black);

  XMapWindow (display, win);
  XSync(display, false);

  // Get min and max for range of color values -- assumed to
  //   be defined by "white", "black"
  min_color = (white > black) ? black : white;
  max_color = (white > black) ? white : black;

  // Pause -- a crude way to reduce the chance that we'll miss
  //   the first few rows of output.
  usleep(2000);  	

  return EXIT_SUCCESS;
}

// wait for user response before ending program
// also allows user to discover coordinates of points
// pre:  display, win are as set up in setup(), other parameters are
//         as described and computed main.
// post:  user has been prompted (see below).
//        mouse clicks in display window cause program to output
//          coordinates to stdout.
//        key press ends function.
void interact(Display * &display, Window &win, 
	      const uint width, const uint height,
	      const double real_min, const double real_max,
	      const double imag_min, const double imag_max) {

  cout << endl << endl;
  cout << "Click on a point in the display to get its coordinates\n";
  cout << "Press any key (with focus in display) to end the program\n";

  //  Choose which events we want to handle  
  XSelectInput(display, win, KeyPressMask | ButtonPressMask);

  // Compute scaling factors (for processing mouse clicks)
  double scale_real = (double) (real_max - real_min) / (double) width;
  double scale_imag = (double) (imag_max - imag_min) / (double) height;

  // Event loop

  XEvent report;
  while (true) {
    
    XNextEvent(display, &report);

    switch ( report.type ) {

    case ButtonPress:

      Window root_return, child_return;
      int root_x_return, root_y_return;
      int win_x_return, win_y_return;
      unsigned int mask_return;
      XQueryPointer(display, win, &root_return, &child_return,
		    &root_x_return, &root_y_return,
		    &win_x_return, &win_y_return,
		    &mask_return);
      cout << "coordinates = ("
	   << real_min + ((double) win_x_return * scale_real)
	   << ", "
	   << imag_min + ((double) (height-1-win_y_return) * scale_imag)
				// height-1-j so y axis displays
				// with larger values at top
	   << ")\n";

      cout.flush();
      break;

    case KeyPress:

      return;
    }
  }
}