package geneFinder;

import ilog.concert.IloException;
import ilog.cplex.IloCplex;

import java.io.*;
import java.util.HashMap;
import java.util.Map;
import java.util.Vector;

import types.*;


/**
* manage the assignment of ambiguous hits to their final position, using a maximum flow formulation
* current methods: GLPK solver, CPLEX linear program
* Copyright (c) 2013,
* Franziska Zickmann, 
* ZickmannF@rki.de, Robert Koch-Institute, Berlin, Germany
* Distributed under the GNU Lesser General Public License, version 3.0
*
*/

public class OptimizeAmbis {

	public static Map<String,Object[]> multiRnas = new HashMap<String,Object[]>();						// store all ambiguous rnas
	public static Map<Gene,Vector<String>> multiRnas_eachCluster = new HashMap<Gene,Vector<String>>();  // stores all variables that belong to one gene
	public static Map<Gene,Object[]> clustNeighbors = new HashMap<Gene,Object[]>();  // first entry: gene vector, second entry: multiRnas, note: if too memory intensive, sacrifice running time and use associated rnas
	
	
	/*
	 * set up the maximum flow formulation
	 */
	
	public static void maxFlowLP(){

		multiRnas.clear();
		multiRnas_eachCluster.clear();
		clustNeighbors.clear();

		long timeBef = System.currentTimeMillis();	

		// the following 3 variables are required for output messages
		int numMulti = 0;	// count all ambiguous reads 
		int numShared = 0;	// count all reads of twin candidate genes
		int processCounter = 0;	// for progress report
		
		int fVar_counter = 0; // necessary to parse cplex solution
		
		StringBuffer binaryVars = new StringBuffer("");			// will include all variables

		boolean noMultis = false;

		try {

			BufferedWriter bw = new BufferedWriter(new FileWriter(GeneFinder.pathOut+"input_it" + GeneFinder.iteration + ".lp"));	// name of lp file

			bw.write("Maximize\n");
			
			StringBuffer targetFunction = new StringBuffer("");		// will include the target function
			
			for(String contigName : GeneFinder.mappedContigs.keySet()){
				
				Contig contig = GeneFinder.mappedContigs.get(contigName);

				for(int posGene = contig.allGenes.size() -1; posGene >= 0; posGene--){

					Gene cluster = contig.allGenes.get(posGene);
					
					boolean hasMultis = false;		// hasMultis and hasTwin indicate if we have to deal with another gene candidate
					boolean hasTwin = false;
					
					cluster.totalCov = 0.0;
					
					if(!FindExonsOfGene.findExonsForGene(cluster)){	// define the exons as they are at the moment, necessary for exon length calculation
						contig.allGenes.remove(cluster);
					}else{
						sumUpExonLengths(cluster, cluster.exonsOfGene);

						if(cluster.twinNode != null){
							// we have to resolve the twin as well					

							cluster.twinNode.totalCov = 0.0;
							if(!FindExonsOfGene.findExonsForGene(cluster.twinNode)){
								cluster.twinNode = null;
							}else{
								sumUpExonLengths(cluster.twinNode, cluster.twinNode.exonsOfGene);
													
								hasTwin = true;
							}
						}


						int multiNumNode = 0;		// count the ambiguous reads of each candidate

						for(String rnaKey : cluster.idTOassociatedRnas.keySet()){

							Rna rna = ((Rna)cluster.idTOassociatedRnas.get(rnaKey)[0]);

							String varName = "x__"+rna.rnaID; 

							if(rna.isMulti == 1 || (rna.isSharedBy.contains(cluster.geneID))){

								if(rna.isMulti == 1){  // to determine the right unique coverage

									multiNumNode++;
									numMulti++;

									cluster.totalCov += (double) (((1.0)/((double) rna.hitNum)) * (double)GeneFinder.readLength);  // ambiguous reads have fewer weight, according to the number of their hits

								}else{
									cluster.totalCov += (double) GeneFinder.readLength;
								}

								if(rna.assignedNum > 1){
									hasMultis = true;

									if(((multiNumNode % 50) == 0) && (targetFunction.length() > 10)){
										bw.write(targetFunction.substring(2));
										targetFunction = new StringBuffer("  ");
									}else{
										targetFunction.append(" + x__"+rna.rnaID + "__" + cluster.geneID + "__f" + "\n");  // include the variable for each readTOgene connection in the target fkt.									
									}
									
									fVar_counter++;
									initialize_MaxFlow(cluster,varName);	// perform the necessary updates of the maps
								}								

							}else{
								cluster.totalCov += (double)GeneFinder.readLength;
							}

						}

						cluster.totalCov = ((double) cluster.totalCov)/((double) cluster.exonLength);

						if(hasMultis){

							cluster.numOfMultis = multiNumNode;
							cluster.uniqueCov = (((double)((cluster.idTOassociatedRnas.keySet().size() - multiNumNode) * GeneFinder.readLength))/((double)cluster.exonLength));

						}

						if(hasTwin){ // now perform the same stuff for the twin candidate

							for(String rnaKey : cluster.twinNode.idTOassociatedRnas.keySet()){

								Rna rna = ((Rna)cluster.twinNode.idTOassociatedRnas.get(rnaKey)[0]);

								String varName = "x__"+rna.rnaID; 

								if(rna.isMulti == 1 || rna.isSharedBy.contains(cluster.twinNode.geneID)){							

									if(!rna.isSharedBy.contains(cluster.twinNode.geneID)){
										// this is non-shared multiple read, so count it
										//cluster.twinNode.numOfMultis++;
										numMulti++;	
									}
									if(rna.isMulti == 1){											
										cluster.twinNode.numOfMultis++;	
										
										cluster.twinNode.totalCov += (double) (((1.0)/((double) rna.hitNum)) * (double)GeneFinder.readLength);
									}

									if(rna.isSharedBy.contains(cluster.twinNode.geneID)){
										numShared++;	// at this point, update the shared rna number								
									}
									
									if(rna.assignedNum > 1){
										
										if(((cluster.twinNode.numOfMultis % 50) == 0) && (targetFunction.length() > 10)){
											bw.write(targetFunction.substring(2));
											targetFunction = new StringBuffer("  ");
										}else{
											targetFunction.append(" + x__"+rna.rnaID + "__" + cluster.twinNode.geneID + "__f" + "\n");
										}
										
										fVar_counter++;
										initialize_MaxFlow(cluster.twinNode,varName);
									}
									

								}else{
									cluster.twinNode.totalCov += (double)GeneFinder.readLength;
								}
							}

							cluster.twinNode.uniqueCov = (((double)((cluster.twinNode.idTOassociatedRnas.keySet().size() - cluster.twinNode.numOfMultis) * GeneFinder.readLength))/((double)cluster.twinNode.exonLength));
							cluster.twinNode.totalCov = ((double) cluster.twinNode.totalCov)/((double) cluster.twinNode.exonLength);
						}

						processCounter++;
						if((processCounter % 2000) == 0){
							
							if((processCounter % 10000) == 0){
								long timeAfterTmp = System.currentTimeMillis();
								if(!GeneFinder.secondPart){
									System.out.println("Processed " + processCounter + " candidates in " + (double) (timeAfterTmp-timeBef)/1000.0 +"s.");
								}
							}
							
							if(!(numMulti == 0 && numShared == 0 &&  multiRnas.keySet().size() == 0)){
								if((targetFunction.length() > 10)){
									bw.write(targetFunction.substring(2));
									targetFunction = new StringBuffer("  ");
								}
							}
						}else{
							if((targetFunction.length() > 10) && (fVar_counter % 1000) == 0){
								bw.write(targetFunction.substring(2));
								targetFunction = new StringBuffer("  ");
							}
						}
					}
				}
			}

			// log messages:
			
			if(!GeneFinder.secondPart){
				System.out.println("Processed " + processCounter + " candidates.");
				System.out.println("Number of multiple rnas: " + numMulti);
				System.out.println("Number of shared rnas: " + numShared);
				System.out.println("Number constraints: " + multiRnas.keySet().size());
				WriteOutput.writeToLogFile("Processed " + processCounter + " candidates. \n" + "Number of multiple rnas: " + numMulti + " \nNumber of shared rnas: " + numShared + " \nNumber constraints: " + multiRnas.keySet().size() + "\n");

				long timeAfter1 = System.currentTimeMillis();
				System.out.println("Time needed for max flow initialization: "+(double) (timeAfter1-timeBef)/1000.0 +"s.");
				WriteOutput.writeToLogFile("Time needed for max flow initialization: "+(double) (timeAfter1-timeBef)/1000.0 +"s.\n");
			}
			
			if(numMulti == 0 && numShared == 0 &&  multiRnas.keySet().size() == 0){
				System.out.println("No multiple rnas, no optimization necessary!");
				WriteOutput.writeToLogFile("No multiple rnas, no optimization necessary! \n");
				noMultis = true;
				GeneFinder.noAmbiOpti = true;
			}else{
				// write out the lp
				if(targetFunction.length() > 3){
					bw.write(targetFunction.substring(2));
					targetFunction = new StringBuffer("");
				}
						
				bw.write("Subject To \n");
				
				// now constraints

				for(String varName : multiRnas.keySet()){

					StringBuffer constraint = new StringBuffer("");			// three different types of necessary constraints
					StringBuffer constraintVar_f = new StringBuffer("");
					StringBuffer constraintVar_Diff = new StringBuffer("");

					for(String var : (Vector<String>) multiRnas.get(varName)[0]){  // set up the different constraints
						
						constraint.append(" + " + var + "\n");
						binaryVars.append(" " + var + "\n");
						constraintVar_f.append(var+"__f >= 0 \n");
						//constraintVar_Diff.append(var+"__f - " + GeneFinder.readLength + " " + var + " <= 0 \n");
						constraintVar_Diff.append(var+"__f - " + var + " <= 0 \n");
					}
					
					constraint.append(" = 1");
					
					bw.write(constraintVar_f.toString() + constraintVar_Diff.toString() + constraint.substring(2) + "\n");

				}

				// now add the weight-constraints

				for(Gene clust : clustNeighbors.keySet()){
					
					double covSum = 0.0;
					double sum_exonL = clust.exonLength;

					StringBuffer constraintF = new StringBuffer("");

					for(Gene neighbor : (Vector<Gene>) clustNeighbors.get(clust)[0]){
						covSum += neighbor.uniqueCov;// + ((1.0/100000.0));   // the ((1.0/100000.0)) ensures that also candidates without unique hits can be chosen
						sum_exonL = sum_exonL + neighbor.exonLength;
					}

					for(String var : (Vector<String>) clustNeighbors.get(clust)[1]){
						constraintF.append(" + " + var + "\n");
					}
					
					// 1) no additional read length factor
					// 2) multiple rnas have a penalty according to their number of ambiguous hits, already included in totalCov
					// 3) competing candidates have a direct influence on the weight due to covSum
					
					double weight = 0; 
					
					covSum = Math.pow((covSum+1),2);
					weight = ((clust.totalCov/covSum) * ((double)clust.exonLength/(double)sum_exonL));
					
					constraintF.append(" <= " + Math.pow((weight+1),2));  // to the power of two leads to a more "the winner takes it all fashion"
					
					bw.write(constraintF.substring(2) + "\n");
					
				}

				//now declare normal variables as binary
			
				bw.write("Binary \n");
				
				int posBin = 0;
				
				for(posBin = 0;posBin<binaryVars.length();){
					int posBin2 = posBin + 500000;
					if(posBin2 < binaryVars.length()){
						bw.write(binaryVars.substring(posBin,posBin2));	// extracted: posBin - (posBin2-1)
					}
					posBin = posBin + 500000;
				}
				
				if((posBin-500000) >= 0){
					bw.write(binaryVars.substring(posBin-500000));
				}
				
				bw.write("END");				
			}
			
			bw.close();

		} catch (IOException e) {
			e.printStackTrace();
		}

		long timeAfter = System.currentTimeMillis();
		if(!GeneFinder.secondPart){
			System.out.println("Time needed to set up lp file: "+ (double) (timeAfter-timeBef)/1000.0 +"s.");
			WriteOutput.writeToLogFile("Time needed to set up lp file: "+ (double) (timeAfter-timeBef)/1000.0 +"s.\n");
			System.out.println();
		}
		
		if(!noMultis){  // now solve the lp
			
			long timeBefSolve = System.currentTimeMillis();
			
			if(GeneFinder.useCPLEX){  // use cplex
				
				if(GeneFinder.optiSolve){
					solveMaxFlowWithCPLEX(numMulti);
					long timeAfterSolve = System.currentTimeMillis();
					System.out.println("Time needed to solve max flow: "+ (double) (timeAfterSolve-timeBefSolve)/1000.0 +"s.");
					WriteOutput.writeToLogFile("Time needed to solve max flow: "+ (double) (timeAfterSolve-timeBefSolve)/1000.0 +"s.\n");
				}else{ // for debugging of human real
					double[] minMax = CalculateScores.assignGeneScores(true);				
					WriteOutput.writeGeneFinderOutput(minMax,"");
				}
				
				// parse solution file
				
				CleanAfterAmbiOpti.parse_solution_and_clean_CPLEX(multiRnas,fVar_counter);

			}else{	// use glpk	
				
				if(GeneFinder.optiSolve){
					solveMaxFlowWithGLPK();
					long timeAfterSolve = System.currentTimeMillis();
					System.out.println("Time needed to solve max flow: "+ (double) (timeAfterSolve-timeBefSolve)/1000.0 +"s.");
					WriteOutput.writeToLogFile("Time needed to solve max flow: "+ (double) (timeAfterSolve-timeBefSolve)/1000.0 +"s.\n");
				}

				// parse solution file
				CleanAfterAmbiOpti.parse_solution_and_clean_GLPK(multiRnas);
			}
		
		}

	}
	
	
	/*
	 * solves MaxFlow with the CPLEX optimizer
	 */
	
	public static void solveMaxFlowWithCPLEX(int numMulti){
		
		try {
			System.out.println("Start cplex solve...");
			IloCplex cplex = new IloCplex();

			cplex.importModel(GeneFinder.pathOut+"input_it" + GeneFinder.iteration + ".lp");
			
			cplex.setParam(IloCplex.IntParam.RootAlg,IloCplex.Algorithm.Network);
			cplex.setParam(IloCplex.DoubleParam.EpGap,0.01);
			cplex.setParam(IloCplex.IntParam.Threads,GeneFinder.numberThreads);
			
			if(numMulti >= 10000000){
				
				if(GeneFinder.memForCplex != -1){
					cplex.setParam(IloCplex.DoubleParam.WorkMem,GeneFinder.memForCplex);
				}
				
				cplex.setParam(IloCplex.StringParam.WorkDir,GeneFinder.pathOut);
				System.out.println("Directory: " + cplex.getParam(IloCplex.StringParam.WorkDir));
				
				cplex.setParam(IloCplex.DoubleParam.PolishTime,1000.0);
				System.out.println("polish time: " + cplex.getParam(IloCplex.DoubleParam.PolishTime));
				
				cplex.setParam(IloCplex.IntParam.HeurFreq,-1);
				
				cplex.setParam(IloCplex.IntParam.NodeFileInd,3);
				System.out.println(cplex.getParam(IloCplex.IntParam.NodeFileInd));
				
				cplex.setParam(IloCplex.BooleanParam.MemoryEmphasis,true);
				System.out.println(cplex.getParam(IloCplex.BooleanParam.MemoryEmphasis));
				
			}
			
			cplex.solve();

			cplex.writeSolution(GeneFinder.pathOut+"solutionCPLEX_it" + GeneFinder.iteration + ".sol");


		} catch (IloException e2) {
			e2.printStackTrace();
		}
		
	}
	
	/*
	 * solves MaxFlow with the glpsol of the GLPK optimizer
	 */
	
	public static void solveMaxFlowWithGLPK(){
		
		try {
			System.out.println("Start glpk solve...");
			
			Runtime rt = Runtime.getRuntime();
			Process firstExe = rt.exec("glpsol --lp " + GeneFinder.pathOut + "input_it" + GeneFinder.iteration + ".lp --output " + GeneFinder.pathOut + "solutionGLPK_out_it" + GeneFinder.iteration + ".out");
			firstExe.waitFor();

		} catch (InterruptedException e) {
			e.printStackTrace();
		} catch (IOException e) {
			e.printStackTrace();
		}
		
	}
	
	
	/*
	 * do the initialization for the max flow optimization
	 * fill the required maps
	 */
	
	public static void initialize_MaxFlow(Gene cluster, String varName){
		
		if(!clustNeighbors.keySet().contains(cluster)){
			Vector<Gene> geneTmp = new Vector<Gene>();	// candidate is not contained so far, so add to map
			Vector<String> varTmp = new Vector<String>();
			varTmp.add(varName+ "__" + cluster.geneID + "__f");
			geneTmp.add(cluster);
			Object[] geneObj = {geneTmp,varTmp};
			clustNeighbors.put(cluster,geneObj);
		}else{
			( (Vector<String>) clustNeighbors.get(cluster)[1]).add(varName+ "__" + cluster.geneID + "__f");  // add this variable to reads of the candidate
		}

		
		if(multiRnas.keySet().contains(varName)){   // add this candidate to the list of genes connected to this read
			((Vector<String>) multiRnas.get(varName)[0]).add(varName+"__"+cluster.geneID);
			
			((Vector<Gene>) multiRnas.get(varName)[1]).add(cluster);
			
			Vector<Gene> clustGenes = (Vector<Gene>) clustNeighbors.get(cluster)[0];
			for(Gene clust : (Vector<Gene>) multiRnas.get(varName)[1]){			// add this candidate to the list of competing genes for all those genes already connected to this read
				if(!clustGenes.contains(clust)){
					( (Vector<Gene>) clustNeighbors.get(cluster)[0]).add(clust);
				}
				
				if(!((Vector<Gene>) clustNeighbors.get(clust)[0]).contains(cluster)){
					((Vector<Gene>) clustNeighbors.get(clust)[0]).add(cluster);
				}
			}
			
		}else{
			Vector<String> vecTmp = new Vector<String>();  // first time ambiguous read was looked at, so add to map
			Vector<Gene> vecGeneTmp = new Vector<Gene>();
			vecGeneTmp.add(cluster);
			vecTmp.add(varName+"__"+cluster.geneID);
			Object[] objTmp = {vecTmp,vecGeneTmp};			
			multiRnas.put(varName,objTmp);															
		}
	}
	
	
	/*
	 * go through all exons and sum up their lengths do determine the total exon length 
	 */
	
	public static void sumUpExonLengths(Gene gene, Vector<int[]> exons){
		
		for(int[] exon :exons){
			gene.exonLength += (exon[1] - exon[0] + 1);
		}
	}
	
}
