// primes_batch - Determine the approximate number of primes within a range.   Requires 
//                at least -np 2 to function correctly.

// improvements to be made:
//	 stop_value, range size and debug as command line arguments 
//	 progress meter
//   echo all command line parameters with debug enabled
//	 hybrid CUDA implementation (one MPI -> N GP-GPU; or one MPI -> N threads -> N GP-GPU)
//	 general fixup/cleanup
//	 VASTART for mpi_error_handler
//   move error handling stuff to an include file
//   add mpi_error_handler to remaining calls, reconcile use of check_mpi with mpi_error_...
//   merge into one source set for all three versions?

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <mpi.h>
#include "check_prime_brute_force.h"

void mpi_error_handler(MPI_Comm *communicator, int *error_code, ...) {
	char error_string[MPI_MAX_ERROR_STRING];
	int my_rank, error_string_length; 
	
	MPI_Comm_rank(*communicator, &my_rank);
	MPI_Error_string(*error_code, error_string, &error_string_length);
	error_string[error_string_length] = '\0'; 
	printf("mpi_error_handler: rank = %d, error_code = %d, error_string = %s\n", 
	  my_rank, *error_code, error_string); 

	exit(1); 
}

int check_mpi(int return_code, char file[128], int line) {
	if (return_code != MPI_SUCCESS) 
		printf("MPI error occurred in %s, on line %d\n", file, line);
	
	return(return_code);
}

int main(int argc, char *argv[]) {
	int rank = 0;
	int size = 0;
	MPI_Status status; 
	MPI_Errhandler errhandler;
	char message[2]; 
	int message_flag;
	int i; 
	
	enum {
		MSG_WORK_REQUEST = 1,
		MSG_WORK_UNIT = 2,
		MSG_WORK_RESULT = 3,
 		MSG_STOP = 4
	};
	
	const int DEBUG = 1;
	const int SERVER = 0;
	
	const int RANGE_START = 0;
	const int RANGE_END = 1;
	
	unsigned long next_value = 3;
	unsigned long stop_value = 1000000000; // must be greater than next_value + range_size
	unsigned long range[2];
	unsigned long range_size = 1000; // dial for the communication to computation ratio
	unsigned long child_result = 0;
	unsigned long total_primes = 2;  // 1 and 2 are prime and we're skipping them
	unsigned long outstanding_ranges = 0; 
	
	double start_time = 0.0;
	double end_time = 0.0;
	
	MPI_Init(&argc, &argv);
	MPI_Comm_rank(MPI_COMM_WORLD, &rank);
	MPI_Comm_size(MPI_COMM_WORLD, &size);
	
	MPI_Comm_create_errhandler(&mpi_error_handler, &errhandler); 
	MPI_Comm_set_errhandler(MPI_COMM_WORLD, errhandler);
  
	if (DEBUG && (rank == 0))
		printf("starting run - number of MPI procesess = %d\n", size); 
	
	// Process command line arguments:
	//   debug
	//   stop_value 
	//   range_size 

	start_time = MPI_Wtime(); 
	
	// server, field requests from clients until we reach the stopping point 
	if (rank == 0) { 

		while (1) { 
			// look for work requests and dispatch next server_value
			message_flag = 0; 
			MPI_Iprobe(MPI_ANY_SOURCE, MSG_WORK_REQUEST, MPI_COMM_WORLD, 
			  &message_flag, &status); 

			if (message_flag) { 
				check_mpi(MPI_Recv(message, sizeof(message), MPI_CHAR, status.MPI_SOURCE, 
				  MSG_WORK_REQUEST, MPI_COMM_WORLD, &status), __FILE__, __LINE__); 

				range[RANGE_START] = next_value; 
				range[RANGE_END] = range[RANGE_START] + range_size; 
				next_value = range[RANGE_END] + 1; 
						         
				MPI_Send(range, 2, MPI_UNSIGNED_LONG, status.MPI_SOURCE, MSG_WORK_UNIT, 
				  MPI_COMM_WORLD); 
				outstanding_ranges++;
			}
	
			// look for work results and update number of primes for this interval
			message_flag = 0; 
			MPI_Iprobe(MPI_ANY_SOURCE, MSG_WORK_RESULT, MPI_COMM_WORLD, &message_flag,
			  &status); 

			if (message_flag) { 
				check_mpi(MPI_Recv(&child_result, 1, MPI_UNSIGNED_LONG, status.MPI_SOURCE, 
				  MSG_WORK_RESULT, MPI_COMM_WORLD, &status), __FILE__, __LINE__); 

				total_primes = total_primes + child_result; 
				outstanding_ranges--;
			} 
			
			// check to see if we're done, if so clean-up and exit
			if (next_value > stop_value) {
			
				// drain all the pending results
				if (DEBUG)
					printf("draining results\n"); 

				while (outstanding_ranges) {
					message_flag = 0; 
					MPI_Iprobe(MPI_ANY_SOURCE, MSG_WORK_RESULT, MPI_COMM_WORLD, 
					  &message_flag, &status); 

					if (message_flag) { 
						check_mpi(MPI_Recv(&child_result, 1, MPI_UNSIGNED_LONG, 
						  status.MPI_SOURCE, MSG_WORK_RESULT, MPI_COMM_WORLD, &status), 
						  __FILE__, __LINE__); 
		
						total_primes = total_primes + child_result;
						outstanding_ranges--; 
					}
				}

				// drain all the work requests 
				if (DEBUG)
					printf("draining work requests\n"); 

				for (i = 1; i < size; i++) {
					message_flag = 0; 
					MPI_Iprobe(MPI_ANY_SOURCE, MSG_WORK_REQUEST, MPI_COMM_WORLD, 
					  &message_flag, &status); 
					if (message_flag) { 
						check_mpi(MPI_Recv(&message, sizeof(message), MPI_CHAR, 
						  status.MPI_SOURCE, MSG_WORK_REQUEST, MPI_COMM_WORLD, &status), 
						  __FILE__, __LINE__); 
					}
				}

				// tell all the children to stop
				if (DEBUG)
					printf("stopping children\n"); 

				for (i = 1; i < size; i++) {
					MPI_Send(&next_value, 1, MPI_UNSIGNED_LONG, i, MSG_STOP, 
					  MPI_COMM_WORLD);
				} 

				break; 
			}

		} // while (1) for server 
	
		end_time = MPI_Wtime(); 
		printf("Roughly %lu primes between 1 and %lu, elapsed time = %8.2f seconds\n", 
		  total_primes, next_value, end_time - start_time); 
	
		MPI_Finalize(); 
		return(0); 		
	}
	
	// client, service requests until we are told to stop
	else { 

		while (1) {
			unsigned long i;  
			unsigned long range[2] = { 0 }; 
			unsigned long total = 0; 

			strcpy(message, "x"); 
			MPI_Send(message, strlen(message) + 1, MPI_CHAR, SERVER, MSG_WORK_REQUEST, 
			  MPI_COMM_WORLD); 
			check_mpi(MPI_Recv(range, 2, MPI_UNSIGNED_LONG, SERVER, MPI_ANY_TAG,
			  MPI_COMM_WORLD, &status), __FILE__, __LINE__);

			switch (status.MPI_TAG) {	
				case MSG_WORK_UNIT: 
					for (i = range[RANGE_START]; i <= range[RANGE_END]; i++) 
						total = total + check_prime_brute_force(i);
					
					MPI_Send(&total, 1, MPI_UNSIGNED_LONG, SERVER, MSG_WORK_RESULT,
					  MPI_COMM_WORLD);
					break; 
	
				case MSG_STOP: 
					MPI_Finalize(); 
					return(0); 
				           
				default:
					// this ought not happen
					break; 
			}	

		} // while (1) 

	} // else for client

} // main