/* Copyright (c) 2012,
 * Mathias Kuhring, KuhringM@rki.de, Robert Koch-Institut, Germany, 
 * All rights reserved. For details, please note the license.txt.
 */

package de.rki.ng4.ipig.mapping;

import java.util.HashMap;
import java.util.Iterator;
import java.util.Vector;
import java.util.Map.Entry;

import de.rki.ng4.ipig.data.Gene;
import de.rki.ng4.ipig.data.GeneSet;
import de.rki.ng4.ipig.data.Peptide;
import de.rki.ng4.ipig.data.PeptideSet;
import de.rki.ng4.ipig.exceptions.ExitException;
import de.rki.ng4.ipig.tools.Configurator;
import de.rki.ng4.ipig.tools.Info;
import de.rki.ng4.ipig.tools.Logger;

/**
 * <p>The Mapper takes a PeptideSet and GeneSet and calculates for each Peptide its occurrences in a linked/referenced Gene
 * ({@link #annotationMapping(PeptideSet, GeneSet)}) or in the all Genes ({@link #alternativeMapping(PeptideSet, GeneSet)}).
 *  
 * <p>Both methods use the Wu-Manber algorithm to match the peptides against amino acid sequences. 
 * Afterwards the cdsMaps are used to map the aa sequence positions to genome positions.</p>
 * 
 * @author Mathias Kuhring
 */
public class Mapper {

	public static String ANNOTATION = "annotation";
	public static String ALTERNATIVE = "alternativ";
	
	private static Position position;

	/**
	 * Given a PeptideSet and a GeneSet, annotationMapping will try to map Peptides to Genes with common UniProtKB-AC, RefSeq or Ensemble_TRS.
	 * 
	 * @param pepset peptides to be mapped 
	 * @param geneset genes to be used as mapping targets
	 * @throws ExitException
	 */
	public static void annotationMapping(PeptideSet pepset, GeneSet geneset) throws ExitException{
		Info info = new Info("map to genome (1)");

		// Peps and genes are linked with the common prots (UniProtKB-AC, RefSeq or Ensemble_TRS)
		HashMap<String, Vector<Integer>> pepmap = pepset.getCommonProts("id");
		HashMap<String, Vector<Integer>> genemap = geneset.getCommonProts();

		Vector<Integer> pepidx;
		Vector<Integer> geneidx;
		Vector<String> pepseqs;

		WuManber wumanber;
		HashMap<Integer, Vector<Integer>> matches;

		Vector<Integer> cds;

		for (Entry<String, Vector<Integer>> entry : pepmap.entrySet()){
			pepidx = entry.getValue();
			geneidx = genemap.get(entry.getKey());

			if (geneidx != null){
				pepseqs = new Vector<String>();
				for (int i : pepidx){
					pepseqs.add(pepset.get(i).getSequence());
				}

				// for each set of peptides with a common protein, the wu-manber alg. is initiated (SHIFT und HASH tables)
				// with the peptides sequences as search pattern
				wumanber = new WuManber(pepseqs, 20);

				for (int i : geneidx){
					Gene gene = geneset.get(i);

					for (int p : pepidx){
						pepset.get(p).setGeneAssigned(true);
						pepset.get(p).addGene(gene.getName());
					}
					
					// each gene with this common protein is used for a search
					matches = wumanber.searchIn(gene.getAaSequence());

					if (!matches.isEmpty()){
						cds = getCdsMap(gene);

						// search matches are evaluated
						for (Entry<Integer, Vector<Integer>> match : matches.entrySet()){
							Peptide pep = pepset.get(pepidx.get(match.getKey()));

							for (int pos : match.getValue()){
								int protStart = getProtStart(pep, gene, pos);
								int protEnd = getProtEnd(pep, gene, pos);

								try{
									getPos(cds, protStart, protEnd);
									write(pep, gene, ANNOTATION);
								}
								catch (IndexOutOfBoundsException e){
									String message = "mapping index error: cds of " + gene.getName() + " might be inconsistent (e.g. check cds coordinates for start and stop codons).";
									Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", message);
								}
							}
						}
					}
					
					Configurator.checkBreak();
				}
			}
		}

		info.stop(pepset.mapPepCount(ANNOTATION), pepset.size());
	}
	
	/**
	 * <p>Given a PeptideSet and a GeneSet, alternativeMapping finds all possible matches for the given Peptides in all given Genes.</p>
	 * 
	 * <p>This method will only involve peptides which are not mapped yet (e.g. with {@link #annotationMapping(PeptideSet, GeneSet)}).</p>
	 * 
	 * @param pepset peptides to be mapped 
	 * @param geneset genes to be used as mapping targets
	 * @throws ExitException
	 */
	public static void alternativeMapping(PeptideSet pepset, GeneSet geneset) throws ExitException{
		Info info = new Info("map to genome (2)");

		// use only unmapped peptides
		Vector<Integer> pepidx = pepset.getUnmapped();
		Vector<String> pepseqs = new Vector<String>();
		for (int i : pepidx){
			pepseqs.add(pepset.get(i).getSequence());
		}

		// initiate wu-manber once with all peptide sequences as search patterns
		WuManber wm = new WuManber(pepseqs, 20);
		HashMap<Integer, Vector<Integer>> matches;
		
		Peptide pep;
		Vector<Integer> cdsMap;
		
		for (Gene gene : geneset.getAll()){
			// use each gene for a the search
			matches = wm.searchIn(gene.getAaSequence());
			
			if (!matches.isEmpty()){
				cdsMap = getCdsMap(gene);
				
				// search matches are evaluated
				for (Entry<Integer, Vector<Integer>> match : matches.entrySet()){
					pep = pepset.get(pepidx.get(match.getKey()));
					pep.setGeneAssigned(true);
					pep.addGene(gene.getName());

					for (int pos : match.getValue()){
						int protStart = getProtStart(pep, gene, pos);
						int protEnd = getProtEnd(pep, gene, pos);
						
						try{
							getPos(cdsMap, protStart, protEnd);
							write(pep, gene, ALTERNATIVE);
						}
						catch (IndexOutOfBoundsException e){
							String message = "mapping index error: cds of " + gene.getName() + " might be inconsistent (e.g. check cds coordinates for start and stop codons).";
							Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", message);
						}
					}
				}
			}

			Configurator.checkBreak();
		}
		
		info.stop(pepset.mapPepCount(ALTERNATIVE), pepset.size());
	}

	// builds a CdsMap, this is vector containing each position in the cds in increasing order
	// e.g. [4,5,6,8,9,13,14,15] where (4,5,6), (8,9) and (13,14,15) are the exons.
	private static Vector<Integer> getCdsMap(Gene gene){
		Vector<Integer> cdsMap = new Vector<Integer>();

		int[] exonStarts = gene.getExonStarts();
		int[] exonEnds = gene.getExonEnds();
		for (int i=0; i<exonStarts.length; i++){
			for (int j=exonStarts[i]; j<exonEnds[i]; j++){
				cdsMap.add(j);
			}
		}

		int test;
		int cdsStart = gene.getCdsStart();
		int cdsEnd =  gene.getCdsEnd();
		for (Iterator<Integer> it = cdsMap.iterator(); it.hasNext();){
			test = it.next();
			if (test < cdsStart) it.remove();
			if (test >= cdsEnd) it.remove();
		}

		return cdsMap;
	}

	// corrects the peptide's start position in the gene's aminoacid sequence (resp. protein) depending on the strand.
	private static int getProtStart(Peptide pep, Gene gene, int pos){
		int protStartPos;

		if (gene.getStrand() == '+'){
			protStartPos = pos;
		}
		else{
			protStartPos = ((gene.getAaSequence().length() + 1) - pos) - pep.getSequence().length();
		}

		return protStartPos * 3;
	}

	// corrects the peptide's end position in the gene's aminoacid sequence (resp. protein) depending on the strand.
	private static int getProtEnd(Peptide pep, Gene gene, int pos){
		int protEndPos;

		if (gene.getStrand() == '+'){
			protEndPos = pos + pep.getSequence().length();
		}
		else{
			protEndPos = (gene.getAaSequence().length() + 1) - pos;
		}

		return protEndPos * 3;
	}

	// uses the peptide's in-protein position to find the in-genome position trough taking the corresponding subset of the cdsMap.
	// testing for gaps in the subset reveals if there are introns resp. if the peptide is located in more then one exon.
	private static void getPos(Vector<Integer> cdsMap, int start, int end) throws IndexOutOfBoundsException{
		Vector<Integer> pepMap = new Vector<Integer>(cdsMap.subList(start, end));

		position = new Position();
		
		position.addStartPos(pepMap.firstElement());
		for (int i=1; i<pepMap.size(); i++){
			if (pepMap.get(i) != pepMap.get(i-1) + 1){
				position.addStartPos(pepMap.get(i));
				position.addEndPos(pepMap.get(i-1)+1);
			}
		}
		position.addEndPos(pepMap.lastElement()+1);
	}

	// this method adds the current calculated Position to the Peptide
	private static boolean write(Peptide pep, Gene gene, String method){
		if ((position.getStartPos().size() > 0) && (position.getEndPos().size() > 0)
				&& (position.getStartPos().size() == position.getEndPos().size())){
			position.setGeneName(gene.getName());
			position.setChrom(gene.getChrom());
			position.setStrand(gene.getStrand());
			
			if (pep.addPosition(position)){
				pep.setMatchMethod(method);
				pep.setGeneMapped(true);
				gene.setUsed(true);
				return true;
			}
			else return false;
		}
		else{
			String message = "mapping error (peptide):\t" + pep.getRow() + "\n" + 
					"mapping error (gene):\t" + gene.toString();
			Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", message);
			return false;
		}
	}

}
