package geneFinder;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.Iterator;
import java.util.TreeMap;
import java.util.Vector;

import types.*;


/**
 * extracts high-coverage clusters as potential genes of prokaryotes
 * 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 ProkaryoteExtraction {

    // variables for log file
	
	public static int numMergedClusters;
	//public static int numFoundNoStart_firstTime;
	//public static int numFoundNoStop_firstTime;
	//public static int numNoFrameFound;
	
	/*
	 * method to analyze the mapping -> extract clusters of high coverage and assign start and stop codons 
	 */
	
	public int initializeClusterSearch(String nameRefFile){

		File refFile;
		if(GeneFinder.useTopHat){
			refFile = new File(nameRefFile+".fa");
		}else{
			refFile = new File(nameRefFile+".fasta");
		}

		int id = 1;

		for(String contigName : GeneFinder.mappedContigs.keySet()){
			Contig thisContig = GeneFinder.mappedContigs.get(contigName);
			StringBuffer contigSeq = new StringBuffer();

			try{ 
				BufferedReader br = new BufferedReader(new FileReader(refFile));
				String line;

				while((line = br.readLine()) != null){
					if(line.startsWith(">")){
						// test if correct contig
						if(line.substring(1).startsWith(contigName)){
							// found right one, now extract sequence
							while(((line = br.readLine()) != null) && (line.length() != 0) &&  (!(line.startsWith(">")))){
								String line2 = "";
								if(Character.isLowerCase(line.charAt(0))){
									for(int i = 0;i<line.length();i++){
										char letter = line.charAt(i);
										letter = Character.toUpperCase(letter);
										line2 += letter;
									}
								}else{
									line2 = line;
								}
								contigSeq.append(line2);
							}
							break;
						}
					}
				}

				if(contigSeq.length() == 0){
					// oops, did not found contig
					System.out.println("Error, could not find contig " + contigName);
					System.exit(0);
				}
		
				// now that we have the sequence, search for areas with high coverage
				Runtime r = Runtime.getRuntime();

				id = searchClusters(thisContig, id, contigSeq);

				double memBef_2 = (r.totalMemory()-r.freeMemory());

				thisContig.positionTOmappingRnas.clear();
				thisContig.positionTOmappingRnas = new TreeMap<Integer,Vector<Rna>>();
				thisContig.positionTOdiff.clear();
				thisContig.positionTOdiff = new TreeMap<Integer,Integer>();
				contigSeq = null;
				r.gc();
				r.gc();

				double memAft_2 = (r.totalMemory()-r.freeMemory());
				if(!GeneFinder.secondPart){
					System.out.println("Memory freed = " + (((memBef_2-memAft_2)/1000.0)/1000.0) + "MB");
					System.out.println();
				}

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

		return id;

	}
	
	/*
	 * searches potential genes on forward and reverse strand, extracts the specific regions respecting reading frame (if possible)
	 */
	
	public int searchClusters(Contig thisContig, int id, StringBuffer contigSeq){

		if(GeneFinder.noAmbiOpti){
			HelperFunctions_GeneSearch.removeAmbiHits(thisContig);
		}
		
		Iterator<Integer> positionIt = thisContig.positionTOmappingRnas.keySet().iterator();
		
		if(GeneFinder.iteration == 2 && !GeneFinder.secondPart){
			System.out.print("Iteration 2. ");
			WriteOutput.writeToLogFile("Iteration 2. ");
		}
		if(!GeneFinder.secondPart){
			System.out.println("Contig: = " + thisContig.contigName);
			WriteOutput.writeToLogFile("Contig: = " + thisContig.contigName + "\n\n");
		}
		
		int startPos = 0;
		int currentPos = 0;
		int nextPos = 0;
		
		// variables for log file
		
		numMergedClusters = 0;

		if(positionIt.hasNext()){
			currentPos = positionIt.next();
		}
		
		// initialize the coverage handling:
		Object[] coverageVecAndPos = new Object[3];
		Vector<Integer> coverageVec = new Vector<Integer>();   // vector is preferred over array because more flexible
		// first the coverage is zero at all positions:
		for(int arrPos=0;arrPos<GeneFinder.readLength;++arrPos){
			coverageVec.add(0);
		}
		
		coverageVecAndPos[0] = coverageVec;
		coverageVecAndPos[1] = currentPos;
		coverageVecAndPos[2] = -1;
		
		if(GeneFinder.endCoverage == -1){
			GeneFinder.endCoverage = (1.0/3.0)*GeneFinder.minCoverage - 0.001; 
			
			if(!GeneFinder.secondPart){
				System.out.println("End coverage estimated from required minimum coverage: " + GeneFinder.endCoverage);
				WriteOutput.writeToLogFile("End coverage estimated from required minimum coverage: " + GeneFinder.endCoverage + "\n\n");
			}
			
		}
		
		if(GeneFinder.maxCov == -1){
			GeneFinder.maxCov = Double.MAX_VALUE;  // so we always accept the coverage
		}

		boolean noMoreCluster = false;
		boolean startedNewCluster = false;   // this boolean ensures that also the last cluster is completed once it has been started due to sufficient coverage (otherwise, if map is empty, cluster is not extracted)
		boolean doNotCountTwice = false;	 // if true we do not perform the coverageVec update for the first current position (when starting a new cluster) because this has already been done with "nextPos"
		
		int numIdentifiedClusters = 0;       // if = 1, this ends the while loop
		
		do{

			Gene cluster = new Gene();
			coverageVecAndPos[2] = -1;
			
			do{
				
				startedNewCluster = false;
			
				int currentCompete = -1;  // stores the current alternative interval, is only -1,-1 if no alternatives exist for current position
				int currentCompeteStart = -1;  // defines the first split determining the currentCompete, necessary to not extract exons, that overlap with currentCompete region
			
				int localExonEnd = -1;
				int endToPassForclosedGenes = -1;  // when we have to close a gene within a currentCompete interval, it is important to grab the last included position as the end
				boolean chooseAlternativeEndInstead = false;
				
				TreeMap<Integer,Vector<Integer>> posiCovMap = new TreeMap<Integer,Vector<Integer>>();  // for each intron, the coverage add after begin and end ist stored
				
				int diff = 0;
				
				if(!doNotCountTwice){

					int covPlus = thisContig.positionTOmappingRnas.get(currentPos).size();
					
					Object[] returnValues = updateCoverageInterval_respectAlternatives(thisContig,covPlus,currentPos,coverageVecAndPos,posiCovMap,-1);
					
					coverageVecAndPos = (Object[]) returnValues[0];
					posiCovMap = (TreeMap<Integer,Vector<Integer>>) returnValues[1];
					
					if(thisContig.positionTOdiff.keySet().contains(currentPos)){
						diff = thisContig.positionTOdiff.get(currentPos); // if there occurred insertions or deletions before this positions add/subtract the difference
					}
					
				}
				
				boolean startWithGene = false;
				if((((Vector<Integer>) coverageVecAndPos[0]).get(0) + diff >= GeneFinder.minCoverage) && (((Vector<Integer>) coverageVecAndPos[0]).get(0) + diff < GeneFinder.maxCov)){
					startWithGene = true;
				}
											
				if(startWithGene){
					startPos = (Integer)coverageVecAndPos[2]; // potential start of a new cluster begins at currentPos - bases covered by already present rnas	
					if((Integer)coverageVecAndPos[2] == -1){
						startPos = currentPos;
					}

					// add also all rnas starting in interval startPos-currentPos

					int pos_temp = (Integer)coverageVecAndPos[2];

					if((Integer)coverageVecAndPos[2] != -1){
						do{							
							//associatedRnas = addRnas(thisContig.positionTOmappingRnas.get(pos_temp),associatedRnas);
							HelperFunctions_GeneSearch.addRnasFromVector(cluster,thisContig.positionTOmappingRnas.get(pos_temp));

							if(thisContig.positionTOmappingRnas.higherKey(pos_temp) != null){
								pos_temp = thisContig.positionTOmappingRnas.higherKey(pos_temp);
							}else{
								pos_temp = currentPos;
							}
						}while((pos_temp != currentPos) && !(pos_temp > currentPos));
					}		


					HelperFunctions_GeneSearch.addRnasFromVector(cluster,thisContig.positionTOmappingRnas.get(currentPos));

					startedNewCluster = true;    

					while(positionIt.hasNext()){ 

						nextPos = positionIt.next();

						int covPlusNext = thisContig.positionTOmappingRnas.get(nextPos).size();

						Object[] returnValues = updateCoverageInterval_respectAlternatives(thisContig,covPlusNext,nextPos,coverageVecAndPos,posiCovMap,-1);

						coverageVecAndPos = (Object[]) returnValues[0];
						posiCovMap = (TreeMap<Integer,Vector<Integer>>) returnValues[1];

						if(thisContig.positionTOdiff.keySet().contains(nextPos)){
							diff = thisContig.positionTOdiff.get(nextPos); // if there occurred insertions or deletions before this positions add/subtract the difference
						}else{
							diff = 0;
						}

						if(((nextPos - currentPos) <= (GeneFinder.readLength)) && (((Vector<Integer>) coverageVecAndPos[0]).get(0) + diff > GeneFinder.endCoverage) && (((Vector<Integer>) coverageVecAndPos[0]).get(0) + diff < GeneFinder.maxCov)){

							endToPassForclosedGenes = nextPos;

							currentPos = nextPos;

							if(currentPos >= currentCompete){
								currentCompete = -1;
								currentCompeteStart = -1;
							}


							HelperFunctions_GeneSearch.addRnasFromVector(cluster,thisContig.positionTOmappingRnas.get(currentPos));


							if(localExonEnd != -1){								
								localExonEnd = currentPos + GeneFinder.readLength;
							}

						}else{
							boolean closeGene = false;

							int basesToAddForOverlap = 0;
							int positionPosiCovMap = HelperFunctions_GeneSearch.findIntronNearNextPos(cluster,nextPos);

							if(posiCovMap.containsKey(positionPosiCovMap)){
								basesToAddForOverlap = posiCovMap.get(positionPosiCovMap).size();
							}

							if((currentCompeteStart != -1 && ((currentCompeteStart - currentPos) < GeneFinder.readLength)) && ((currentCompete + basesToAddForOverlap + 1) >= nextPos) && (((Vector<Integer>) coverageVecAndPos[0]).get(0) + diff < GeneFinder.maxCov)){
								// +1 because currentCompete defines the intron end
								// have to go on because we are still within the cluster, so just decide if this part shall be included or not								

								int goOn = 0;

								if((nextPos - currentPos) > (GeneFinder.readLength)){
									goOn = -1;
								}

								if(goOn != -1){ 

									if(goOn == 0){									
										HelperFunctions_GeneSearch.addRnasFromVector(cluster,thisContig.positionTOmappingRnas.get(nextPos));
									}									

									currentPos = nextPos;							

									if(currentPos >= currentCompete){
										currentCompete = -1;
										currentCompeteStart = -1;
									}
								}else{

									currentCompete = -1;
									currentCompeteStart = -1;
									if(endToPassForclosedGenes != -1){
										chooseAlternativeEndInstead = true;
									}
									closeGene = true;
								}

							}else{
								closeGene = true;
							}


							if(closeGene){

								int basesToAdd =  GeneFinder.readLength;

								if(chooseAlternativeEndInstead){
									currentPos = endToPassForclosedGenes;
								} 

								if(((currentCompeteStart != -1 && ((currentCompeteStart - currentPos) < GeneFinder.readLength))) && currentCompete > currentPos){

									currentPos = currentCompete + 1;   // nextPos is bigger than split end, but we have to consider the split anyway

									// look new number up in posiCovMap

									if(posiCovMap.containsKey(currentCompete)){
										basesToAdd = posiCovMap.get(currentCompete).size();
									}

								}

								// extract this cluster, interval [startPos,currentPos+readLength]							

								id = clustIni(cluster, thisContig, contigSeq, startPos, id, currentPos,basesToAdd);

								numIdentifiedClusters++;
								currentPos = nextPos;  // now nextPos is potential new start								

								doNotCountTwice = true;
								break;
							}
						}							
					}

				}else{
					if(positionIt.hasNext()){
						currentPos = positionIt.next();
						doNotCountTwice = false;
						
						cluster.idTOassociatedRnas.clear();
						cluster.possibleIntrons.clear();
												
					}else{
						break;
					}
				}

				if(!positionIt.hasNext() && (numIdentifiedClusters < 1) && startedNewCluster){  // to grab very last cluster
										
					int basesToAdd =  GeneFinder.readLength;
				
					if(chooseAlternativeEndInstead){
						currentPos = endToPassForclosedGenes;
					} 
					
					if(((currentCompeteStart != -1 && ((currentCompeteStart - currentPos) < GeneFinder.readLength))) && currentCompete > currentPos){
						
						currentPos = currentCompete + 1;   // nextPos is bigger than split end, but we have to consider the split anyway
						
						// look new number up in posiCovMap
						
						if(posiCovMap.containsKey(currentCompete)){
							basesToAdd = posiCovMap.get(currentCompete).size();
						}
					}
					
					id = clustIni(cluster, thisContig, contigSeq, startPos, id, currentPos,basesToAdd);
	
					numIdentifiedClusters++;
				}

			}while(positionIt.hasNext() && (numIdentifiedClusters < 1));

			if(!positionIt.hasNext() && (numIdentifiedClusters < 1)){
				// reached end of reference sequence, no cluster has been extracted this time, so stop
				noMoreCluster = true;
				break;
			}

		
			// now test for overlap merging:

			if(thisContig.allGenes.size() > 0){

				Gene clusterBeforeInVec = thisContig.allGenes.get(thisContig.allGenes.size()-1);

				if(MergeClusters.checkIfNeedMerge(cluster, clusterBeforeInVec,thisContig,contigSeq)){
					numMergedClusters++;
				}else{
					thisContig.allGenes.add(cluster);
				}
								
			}else{
				thisContig.allGenes.add(cluster);
			}
		

			if(thisContig.allGenes.size() > 1){

				Gene clustBef = thisContig.allGenes.get(thisContig.allGenes.size()-2);
				if(clustBef.operonOrfs.isEmpty()){
					Object[] returnArr = Prokaryote_Specials.define_OrfsInOperon(clustBef.sequence,clustBef);
					clustBef.operonOrfs = (Vector<int[]>) returnArr[0];  // first entry is always 1-d array indicating whether normal gene sequence follows (1) or operon sequence (-1)
					clustBef.operonDirectionIsForward = (Boolean) returnArr[1];
					if(clustBef.operonDirectionIsForward){
						clustBef.onRevStrand = false;
					}else{
						clustBef.onRevStrand = true;
					}
					
				}			

			}
		
			if(!positionIt.hasNext()){
				noMoreCluster = true;			
				break;
			}

			numIdentifiedClusters--;

		}while(!noMoreCluster);

		// check for the last extracted cluster if it has a proper stop or not and further check for twins

		if(thisContig.allGenes.size() != 0){
			Gene clustBef = thisContig.allGenes.get(thisContig.allGenes.size()-1);

			if(clustBef.operonOrfs.isEmpty()){
				Object[] returnArr = Prokaryote_Specials.define_OrfsInOperon(clustBef.sequence, clustBef);
				clustBef.operonOrfs = (Vector<int[]>) returnArr[0]; // first entry is always 1-d array indicating whether normal gene sequence follows (1) or operon sequence (-1)
				clustBef.operonDirectionIsForward = (Boolean) returnArr[1];
				if(clustBef.operonDirectionIsForward){
					clustBef.onRevStrand = false;
				}else{
					clustBef.onRevStrand = true;
				}
			}

		}	

		if(!GeneFinder.secondPart){
			String s = "";
			s += "No more clusters can be found \n";
			s += "Total identified clusters: " + thisContig.allGenes.size() + "\n";
			s += "Number mergings: " + numMergedClusters + "\n";
			
			WriteOutput.writeToLogFile(s + "\n");
			System.out.println(s);
		}

		return id;
	}
	
	/*
	 * check if only multiRnaSupport
	 */
	
	public static boolean checkIfOnlyMultiRnas(Gene gene){
		
		for(String rnaString : gene.idTOassociatedRnas.keySet()){
			Rna rna = ((Rna)gene.idTOassociatedRnas.get(rnaString)[0]);
			if(rna.isMulti == 0){
				return true;
			}
		}
		return false;
	}
	
	/*
	 * for all rnas mapping to this gene, remove the gene alignment
	 */
	
	public static void removeGenesRnas(Gene gene, Contig thisContig){
		
		for(String rnaID : gene.idTOassociatedRnas.keySet()){
			Rna rna = (Rna) gene.idTOassociatedRnas.get(rnaID)[0];
			
			for(Object[] info : rna.contigsMappedOn){
				
				int alignPos = ((Integer) info[1]).intValue();
				
				if((((Contig) info[0]).equals(thisContig)) && (alignPos > gene.startPos) && (alignPos < (gene.stopPos-GeneFinder.readLength))){				
					rna.contigsMappedOn.removeElement(info);
					rna.assignedNum = rna.assignedNum -1;
					break;				
				}
			}

		}
	}

	/*
	 * updates the coverage entries in the vector representing the current covered interval
	 */
	
	public static Object[] updateCoverageInterval_respectAlternatives(Contig thisContig, int covPlus, int currentPos, Object[] coverageVecAndPos, TreeMap<Integer,Vector<Integer>> posiCovMap,int considerSpliceSite){
		
		boolean addedValues = false;
		
		int splitDiff = 0;    // necessary to consider splice site, do not update with covplus if we exceed splice site	
		if(considerSpliceSite != -1){
			splitDiff = (considerSpliceSite-currentPos);
		}
		
		Vector<Integer> covVecClone = new Vector<Integer>();    // stores all coverage add values derived by posiCovMap
		Object[] returnObject = HelperFunctions_GeneSearch.lookIntoPosiCovMap(posiCovMap, currentPos);

		addedValues = (Boolean) returnObject[0];
		covVecClone = (Vector<Integer>) returnObject[1];
		posiCovMap = (TreeMap<Integer,Vector<Integer>>) returnObject[2];
		
		if(currentPos - (Integer) coverageVecAndPos[1] > GeneFinder.readLength){

			Vector<Integer> coverageVecNew = new Vector<Integer>();
			// initialize:
			for(int arrPos=0;arrPos<GeneFinder.readLength;++arrPos){
				if((considerSpliceSite != -1) && (arrPos >= splitDiff)){
					coverageVecNew.add(0);
				}else{
					coverageVecNew.add(covPlus);
				}
			}
			
			coverageVecAndPos[0] = coverageVecNew;
			coverageVecAndPos[2] = -1;
			

		}else{
			
			Vector<Integer> covVec = (Vector<Integer>) coverageVecAndPos[0];
			for(int pos = 0;pos<(currentPos- (Integer) coverageVecAndPos[1]);++pos){
				covVec.remove(covVec.firstElement());
				covVec.add(0);
			}
			for(int vecPos=0;vecPos < covVec.size();++vecPos){
				if(!((considerSpliceSite != -1) && (vecPos >= splitDiff))){
					covVec.set(vecPos,(covVec.get(vecPos) + covPlus));    // correct?
				}
				
			}
			
			if((((Integer)coverageVecAndPos[2] == -1)) || (currentPos-((Integer) coverageVecAndPos[2]) > (GeneFinder.readLength))){
				// update overlap-position
				boolean foundPos = false;
				int pos_temp = (Integer) coverageVecAndPos[2];
				if(pos_temp == -1){
					coverageVecAndPos[2] = (Integer) coverageVecAndPos[1];
				}else{
					if(!addedValues){  // if we have introns, then stay with the current start position
						do{	
							if(thisContig.positionTOmappingRnas.lastKey() != pos_temp){
								pos_temp = thisContig.positionTOmappingRnas.higherKey(pos_temp);
							}else{
								pos_temp = (Integer) coverageVecAndPos[1];
							}
							if((currentPos-(pos_temp) <= (GeneFinder.readLength))){
								coverageVecAndPos[2] = pos_temp;
								foundPos = true;
							}
						}while(!foundPos || !(pos_temp >= (Integer) coverageVecAndPos[1]));
					}
				}
			}
		}
		
		coverageVecAndPos[1] = currentPos;
		
		if(addedValues){
			for(int posVec = 0; posVec < GeneFinder.readLength; ++posVec){
				((Vector<Integer>) coverageVecAndPos[0]).setElementAt((((Vector<Integer>)coverageVecAndPos[0]).get(posVec) + covVecClone.get(posVec)),posVec);
			}
		}
		
		Object[] toReturn = {coverageVecAndPos,posiCovMap,addedValues};
		return toReturn;
	}
	
	/*
	 * first extraction of each cluster, simply regard the high coverage interval
	 */
	
	public static int clustIni(Gene cluster, Contig thisContig, StringBuffer contigSeq, int startPos, int id, int currentPos, int basesToAdd){
		
		basesToAdd = basesToAdd + GeneFinder.readLength;
		
		startPos = (int)Math.max(0,(startPos-GeneFinder.readLength));
		
		cluster.sequence = contigSeq.substring(startPos,Math.min(currentPos+basesToAdd,contigSeq.length()));

		cluster.geneID = id++;
		cluster.onRevStrand = false;
		cluster.startPos = startPos;		
		cluster.stopPos = Math.min(currentPos+basesToAdd-1,contigSeq.length()-2); // note: without "-1" stopPos would be the position necessary for seq-extraction = actual stop + 1 
		
		cluster.coreSeq = (cluster.stopPos-cluster.startPos);
		cluster.realDirectionNotKnown = true;	
		
		return id;
	}
	
}


