package geneFinder;

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_GLPK {


		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();
				
				// 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 glpsol of the GLPK optimizer
		 */
		
		public static void solveMaxFlowWithGLPK(){		
		
			System.out.println("Start glpk solve...");
				
			String firstExe = "python " + GeneFinder.pathToHelpFiles+"lpCall.py " + GeneFinder.pathOut + "input_it" + GeneFinder.iteration + ".lp " + GeneFinder.pathOut + "solutionGLPK_out_it" + GeneFinder.iteration + ".out"; 	
			Giira.callAndHandleOutput(firstExe,true);
			
		}
		
		
		/*
		 * 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);
			}
		}
		
	}

