package be.ac.vub.ir.mpi;

import be.ac.vub.ir.util.Chrono;
import be.ac.vub.ir.util.GetOpt;
import mpi.MPI;
import mpi.MPIException;
import mpi.Status;

/**
 * @author Joris Borms
 * parallel and sequential implementation of a prime number counter
 */

public class CountPrimes {
		
	// mpi user-defined tags
	private final static int INPUT_TAG = 1;
	private final static int RESULT_TAG = 2;
	
	// default program parameters
	final static int DEFAULTSIZE = 100000;

	public static void main(String[] args) throws MPIException {		
		
		GetOpt options = new GetOpt("n:",args); // allows you to parse program options. 
		// ':' means that an argument is following after the option (here: -n 1000)
	
		MPI.Init(args);
		int rank = MPI.COMM_WORLD.Rank();
		int size = MPI.COMM_WORLD.Size();
		int nbrSlaves = size - 1;
	
		if (rank == 0) { // we choose rank 0 for master program
			
			// initialise data
			String ns = options.getOptionParam('n');
			int arraySize = (ns != null) ? Integer.parseInt(ns) : DEFAULTSIZE ;

			int[] data = createAndFillArray(arraySize);
			
			// *** PARALLEL RUN ***
			Chrono parChrono = new Chrono();			
			
			// divide data over slaves
			int slavedata = arraySize / nbrSlaves; // # data for one slave
			int rest = arraySize - (slavedata * nbrSlaves);
			int index = 0;

			for (int slaveID=1; slaveID < size; slaveID++) {	
				MPI_sendIntArr(data, index, slavedata + rest, slaveID, INPUT_TAG);
				index += slavedata + rest;
				rest = 0; // first slave gets a bit more if n is not divisable by slaves
			}			
			
			// slaves are working...
			
			int nbrPrimes = 0;
			for (int slaveID=1; slaveID < size; slaveID++) 
				nbrPrimes += MPI_recvInt(slaveID, RESULT_TAG);
			
			long partime = parChrono.stop();
			System.out.println("Parallel solution found "+nbrPrimes+" primes in "+arraySize+" integers.");
			
			// *** SEQUENTIAL RUN (master only) ***
			Chrono seqChrono = new Chrono();
			int nbrPrimesSeq = countPrimes(data);
			long seqtime = seqChrono.stop();
			
			if (nbrPrimesSeq == nbrPrimes)
				System.out.println("Sequential solution found the same number of primes. OK.");
			else
				System.err.println("Sequential solution found "+nbrPrimesSeq+" primes which is different!!");
			
			System.out.println();
			System.out.println("PARALLEL RUNTIME   = "+parChrono);
			System.out.println("SEQUENTIAL RUNTIME = "+seqChrono);
			System.out.println("speedup            = "+((float) seqtime/partime));
			System.out.println("efficiency         = "+((float) seqtime/partime/nbrSlaves));
			System.out.println(" ++++++++++ finished +++++++++++ ");
		}
		else {
			// *** Slave Program ***
			Chrono slaveChrono = new Chrono();
			int[] array = MPI_recvIntArr(0, INPUT_TAG);	
			
			Chrono countingChrono = new Chrono();
			int result = countPrimes(array);	 // calculating...
			countingChrono.stop();
			
			MPI_sendInt(result, 0, RESULT_TAG);
			slaveChrono.stop();
			// prints after sending back result
			System.out.println("slave "+rank+" received "+array.length+" integers and found "+result+" primes (total time: "+slaveChrono+", computing: "+countingChrono+")");		
		}
		
		MPI.Finalize(); // Don't forget!!
	}
	// *** Array Filling ***
	static int[] createAndFillArray(int n){ // fills with all odd numbers
		int[] data = new int[n];
		for (int i=0;i<n;i++)
			data[i] = 1 + 2*i;;
		return data;
	}
	
	// *** Sequential Prime Counting ***
	public static int countPrimes(int[] data) {
		int np = 0;
		for (int value: data)
			if (isPrime(value))
				np++;
		return np;		
	}
	
	public static boolean isPrime(int value){ // this is not the most efficient implementation, we know...
		for (int j=2 ; j<value ; j++){
			if ((value % j) == 0)
				return false;
		}
		return true;
	}
	// *** Utility Communication Functions ***
	// custom blocking send/receive for integers & integer arrays	
	public static void MPI_sendInt(int arg, int dest, int tag) throws MPIException {
		int[] buff = new int[] {arg};
		MPI.COMM_WORLD.Send(buff, 0, 1, MPI.INT, dest, tag);
	}
	
	public static int MPI_recvInt(int source, int tag) throws MPIException {
		int buff[] = new int[1]; // allocate buffer size of 1
		MPI.COMM_WORLD.Recv(buff, 0, 1, MPI.INT, source, tag);
		return buff[0];
	}
	
	public static void MPI_sendIntArr(int[] array,int dest,int tag) throws MPIException {
		MPI_sendIntArr(array,0,array.length,dest,tag);
	}
	
	public static void MPI_sendIntArr(int[] array,int offset,int count,int dest,int tag) throws MPIException {
		MPI.COMM_WORLD.Send(array, offset, count, MPI.INT, dest, tag);
	}
	
	public static int[] MPI_recvIntArr(int source, int tag) throws MPIException {
		// to allocate the array we first have to probe for the message size (count)
		Status status = MPI.COMM_WORLD.Probe(source, tag);
		int[] array = new int[status.count];
		MPI.COMM_WORLD.Recv(array, 0, status.count, MPI.INT, source, tag);
		return array;		
	}
}