package exp.twostep;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileWriter;
import java.io.IOException;
import java.net.InetAddress;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Random;
import java.util.stream.Collectors;

import mpi.*;

/**
 * This class is almost the same as OpenMPIstep2.
 * But this class shuffles for 10000 times;
 * */
public class OpenMPIstep2p {
	
	public final static int tag = 233333; //distribute data
	public static final int source = 0;//rank 0 is the controller.'
	
	
	public static boolean allFalse(boolean[] a){
		boolean f = true;
		for (int i = 1; i < a.length; i++){
			f = f & !a[i];
		}
		return f;
	}
	
	/**
	 * argument 
	 * 1. mrna, 2.met, 3.tf, 4.cnv, 5.met/tf as condition, 6. thread, 7. shuffleFile; 8. number of shuffles (optional)
	 * */
	static String computerName = "Failed";
	public static void main(String[] args) throws MPIException {
		MPI.Init(args);
		MPI.COMM_WORLD.setErrhandler(MPI.ERRORS_RETURN);
		System.out.println(System.getenv().get("java.library.path"));
		System.out.println(System.getProperty("java.library.path"));
		System.out.println("args length = " + args.length);
		if (args.length >= 8)
			System.setProperty("java.library.path", System.getProperty("java.library.path") + File.pathSeparator + args[7]);
		System.out.println("2 : " + System.getProperty("java.library.path"));
		System.out.println("MPI INIT");
		//test library first
		MI mi = new MI();
		
//		double[][] a1 = new double[1][20];
//		double[][] b1 = new double[1][20];
//		double[][] c1 = new double[1][20];
//		double d = mi.partialMutualInformation(a1, b1, c1, 4); //use this to trigger the error
		
		long time1 = System.currentTimeMillis();
		
		final int repeat = 10000;//final
		int myrank = MPI.COMM_WORLD.getRank();
		int size = MPI.COMM_WORLD.getSize();
		
		try{
			computerName = InetAddress.getLocalHost().getHostName();
		}catch (Exception e){
			System.err.println("Had problem fetching computer name! Now exit.");
			System.exit(-2);
		}
		
		System.out.println("MPI argument initialized");
		
		if (myrank == 0){
			System.out.println("(0)rank 0 started.");
			Loading l = new Loading(args);
			System.out.println("Loading finished");
			List<String> tfList = l.getTFList();
			ArrayList<Pair<String,String> > geneMet = l.getGeneMetCombination();
			DataRepresent met = l.getMetData();
			DataRepresent mrna = l.getmRNAData();
			DataRepresent cnv = l.getCNVData();
			int mode = l.getMode();
			int cores = l.getThread();
			String[] extraArgument = l.getExtraArgument();
			double[][] position = l.getLocation();
			int[] sendBuffer = new int[2];//indicate length
			sendBuffer[0] = l.getLength();
			sendBuffer[1] = repeat;
			MPI.COMM_WORLD.bcast(sendBuffer, 2, MPI.INT, 0);//send size and repeat
			int rowNum = tfList.size();
			int colNum = geneMet.size();
			boolean[] status = new boolean[size];
			Arrays.fill(status, true);
			status[0] = false;
			int sendTo = 1;
			int i = 0;//current work
			int readCount = 0;
			//specific job deploy
			Map<Integer, Double> resultS = new HashMap<Integer, Double>();
			for (int h = 0; h < position.length; h++){
				resultS.put((int)position[h][0], 999.0);
			}
			int recvPos = 0;
			while (i < position.length){
				double[] sendBuffer2 = new double[2 + l.getLength() * 4];
				System.arraycopy(position[i], 0, sendBuffer2, 0, 2);//sendBuffer2 location, cMI, 
				double[] Met, Tf, Cnv, Gene;
				Cnv = new double[l.getLength()];
				Arrays.fill(Cnv, 0);
				int tfPos = (int)position[i][0] % rowNum;
				int metgenePos = (int)position[i][0] / colNum;
				Tf = mrna.content.get(tfList.get((tfPos)));
				Gene = mrna.content.get(geneMet.get(metgenePos).first);
				Met = met.content.get(geneMet.get(metgenePos).second);
				if (cnv.content.containsKey(geneMet.get(metgenePos).first)){
					Cnv = cnv.content.get(geneMet.get(metgenePos).first);
				}
				System.arraycopy(Gene, 0, sendBuffer2, 2, l.getLength());
				System.arraycopy(Cnv, 0, sendBuffer2, 2 + 2 * l.getLength(), l.getLength());
				if (mode == Worker.MET){
					System.arraycopy(Tf, 0, sendBuffer2, 2 + 1 * l.getLength(), l.getLength());
					System.arraycopy(Met, 0, sendBuffer2, 2 + 3 * l.getLength(), l.getLength());
				}
				if (mode == Worker.TF){
					System.arraycopy(Met, 0, sendBuffer2, 2 + 1 * l.getLength(), l.getLength());
					System.arraycopy(Tf, 0, sendBuffer2, 2 + 3 * l.getLength(), l.getLength());
				}
				
				for (sendTo = 1; sendTo < size; sendTo++){
					if (status[sendTo]){
						MPI.COMM_WORLD.send(sendBuffer2, sendBuffer2.length, MPI.DOUBLE, sendTo, tag);
						System.out.printf("(rank 0 on %s) send data(ID %d) to (rank %d)\n", computerName, (int)sendBuffer2[0], sendTo);
						status[sendTo] = false;
						i++;
						break;//do not repeat!!!
					}
				}
				double[] recvBuff = new double[3];
				if (allFalse(status)){
					Status s = MPI.COMM_WORLD.recv(recvBuff, 3, MPI.DOUBLE, MPI.ANY_SOURCE, tag);
					System.out.printf("(rank %d on %s) receiving from (rank %d) ID=%d\n", myrank, computerName, s.getSource(), (int)recvBuff[0]);
					readCount++;// keep a record of how many result are returned
					status[s.getSource()] = true;
					resultS.replace((int)recvBuff[0], recvBuff[1]);
				}

			}
			// When previous loop exited, we should check the rest threads
			// possible status: status[1:end] has 1 true or has 0 true
			double[] recvBuff = new double[3];
			for (int j = 1; j < status.length; j++){
				if (status[j] == true)
					continue;
				Status ds = MPI.COMM_WORLD.recv(recvBuff, 3, MPI.DOUBLE, j, tag);//wait for those who are not able to return values.
				System.out.printf("(rank %d on %s) receiving from (rank %d) ID=%d\n", myrank, computerName, ds.getSource(), (int)recvBuff[0]);
				resultS.replace((int)recvBuff[0], recvBuff[1]);
			}//finish the rest 
			double[] endSig = new double[2 + 4 * l.getLength()];
			Arrays.fill(endSig, 0);
			endSig[0] = -1;
			for (int k = 1; k < size; k++)//send stop signal
				MPI.COMM_WORLD.send(endSig, endSig.length, MPI.DOUBLE, k, tag);// This signal tells all other process to stop receiving
			
//			MPI.COMM_WORLD.barrier();
			
//			recvBuff = new double[3];
//			for (int k = 0; k < position.length; k++){
//				MPI.COMM_WORLD.recv(recvBuff, 2, MPI.DOUBLE, MPI.ANY_SOURCE, tag);
//				System.arraycopy(recvBuff, 0, resultS[k], 0, 2);
//			}
			
			List<String> fl = resultS.entrySet().stream().sorted((Entry<Integer, Double> x, Entry<Integer, Double> y)->{
				if (x.getKey() > y.getKey()){
					return 1;
				}else if (x.getKey() == y.getKey()){
					return 0;
				}else{
					return -1;
				}
			}).map(x -> String.format("%d\t%.8f\n", (int)(x.getKey()), x.getValue())).collect(Collectors.toList());
			// was Math.round(x.getKey()) would convert x.getKey() to float !!!!
			
			try (FileWriter fw = new FileWriter("ID_sig.txt");
				BufferedWriter bw = new BufferedWriter(fw)){
				for (i = 0; i < fl.size(); i++){
					bw.write(fl.get(i));
				}
			}catch (FileNotFoundException fnfe){
				fnfe.printStackTrace();
			}catch (IOException ioe){
				ioe.printStackTrace();
			}
			try (FileWriter fw = new FileWriter("output_sig.txt");
				BufferedWriter bw = new BufferedWriter(fw)){
				for (i = 0; i < fl.size(); i++){
					bw.write(fl.get(i));
				}
			}catch (FileNotFoundException fnfe){
				fnfe.printStackTrace();
			}catch (IOException ioe){
				ioe.printStackTrace();
			}			
			
			long time2 = System.currentTimeMillis();
			System.out.println("Finished");
			System.out.printf("Consumed time = %f s.\n", (time2 - time1) / 1000.0d);
			
			
			
			
		}else{
			System.out.printf("rank %d started\n", myrank);
			int[] recvBuffer = new int[2];
			MPI.COMM_WORLD.bcast(recvBuffer, 2, MPI.INT, 0);
			System.out.printf("(rank %d) length broadcast: %d, repeat: %d\n", myrank, recvBuffer[0], recvBuffer[1]);
			
			final int length = recvBuffer[0];
			final int repeat1 = recvBuffer[1];
			
			receiveAndProcess(length, repeat1, myrank);
			
		}
		
		MPI.Finalize();
	}
	
	public static void receiveAndProcess(int length, int repeat, int myrank) throws MPIException{
		int len = 2 + 4 * length;//id cmi
		double[] recvBuffer = new double[len];
		ArrayList<double[]> data = new ArrayList<double[]>();//original data
		ArrayList<double[]> result = new ArrayList<double[]>();
		while (true){
			MPI.COMM_WORLD.recv(recvBuffer, len, MPI.DOUBLE, 0, tag);
			System.out.printf("(rank %d on %s) receiving ID = %d\n", myrank, computerName, (int)recvBuffer[0]);
			if (recvBuffer[0] < 0){
				System.out.printf("(rank %d on %s) exiting\n", myrank, computerName);
				break;
			}
			MI mi = new MI();
			int count = 0;
			double[][] a = new double[1][length];
			double[][] b = new double[1][length];
			double[][] c = new double[2][length];
			int k = (int)Math.ceil(0.16 * length);				
			double[] d = recvBuffer;
			System.arraycopy(d, 2, a[0], 0, length);
			System.arraycopy(d, 2 + length, b[0], 0, length);
			System.arraycopy(d, 2 + 2 * length, c[0], 0, length);
			System.arraycopy(d, 2 + 3 * length, c[1], 0, length);
			for (int i = 0; i < repeat; i++){
				rearrange(c[1]);
				count += MI.partialMutualInformation(a, b, c, k) > d[1] ? 1 : 0; /////SHUFFLE
			}
			double[] buffer = new double[3];
			buffer[0] = d[0];
			buffer[1] = 1.0d * count / repeat;
			buffer[2] = myrank;
			MPI.COMM_WORLD.send(buffer, 3, MPI.DOUBLE, 0, tag);
	
		}
//		System.out.printf("(rank %d) received %d data\n", myrank, data.size());
		
		//then calculation
//		for (int i = 0; i < data.size(); i++){
//			MI mi = new MI();
//			int count = 0;
//			double[][] a = new double[1][length];
//			double[][] b = new double[1][length];
//			double[][] c = new double[2][length];
//			int k = (int)Math.ceil(0.16 * length);
//			double[] d = data.get(i);
//			System.arraycopy(d, 2, a[0], 0, length);
//			System.arraycopy(d, 2 + length, b[0], 0, length);
//			System.arraycopy(d, 2 + 2 * length, c[0], 0, length);
//			System.arraycopy(d, 2 + 3 * length, c[1], 0, length);
//			
//			for (int j = 0; j < repeat; j++){
//				rearrange(c[1]);
//				count += MI.partialMutualInformation(a, b, c, k) > d[1] ? 1 : 0; /////SHUFFLEŁĄŁĄŁĄŁĄŁĄŁĄ
//				
//			}
//			
//			double[] pp = new double[]{d[0], 1.0d * count / repeat};
//			result.add(pp.clone());
//			if (i % 100 == 0){
//				System.out.printf("(rank %d) processed %d line.\n", myrank, i);
//			}
//		}
//		MPI.COMM_WORLD.barrier();//sync
//		for (int i = 0; i < result.size(); i++){
//			double[] buffer = new double[2];
//			System.arraycopy(result.get(i), 0, buffer, 0, 2);//ID and pvalue copied
//			MPI.COMM_WORLD.send(buffer, 2, MPI.DOUBLE, 0, tag);
//		}
		
		
		
	}
	
	public static void rearrange(double[] d){
		int length = d.length;
		Random r = new Random();
		for (int i = 0; i < d.length; i++){
			int pos = r.nextInt(length - i) + i;
			double s = d[i];
			d[i] = d[pos];
			d[pos] = s;
		}
	}
}
