/* 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.data;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Vector;

import de.rki.ng4.ipig.exceptions.ExitException;
import de.rki.ng4.ipig.exceptions.FormatException;
import de.rki.ng4.ipig.mapping.Mapper;
import de.rki.ng4.ipig.mapping.Position;
import de.rki.ng4.ipig.tools.Configurator;
import de.rki.ng4.ipig.tools.Info;
import de.rki.ng4.ipig.tools.Logger;

/**
 * PeptideSet handles the peptides that should be mapped to a genome.
 * This includes import methods and creation of a Peptide object for each peptide imported from a file.
 * 
 * @author Mathias Kuhring
 *
 */
public class PeptideSet {

	private Vector<Peptide> peptides;
	private String header = "prot_acc	prot_desc	pep_query	pep_isunique	pep_exp_z	pep_score	pep_seq	pep_var_mod	pep_var_mod_pos";

	/**
	 * <p>Initializing a PeptideSet by giving a file (txt or mzid) of peptides identified from mass spectra.</p>
	 * 
	 * <p>A txt file must be tab separated and start with the following header:<br>
	 * "prot_acc	prot_desc	pep_query	pep_isunique	pep_exp_z	pep_score	pep_seq	pep_var_mod	pep_var_mod_pos"</p>
	 * 
	 * <p>A mzid file must be a valid mzIdentML format file in version 1.0.0 
	 * (<a href="http://www.psidev.info/index.php?q=node/403">Link</a>). 
	 * It will be checked with a XML schema definition.</p>
	 * 
	 * @param filename
	 * @throws ExitException
	 */
	public PeptideSet(String filename) throws ExitException{
		peptides = new Vector<Peptide>();
		readPeptides(filename);
	}

	/**
	 * Returns the number of Peptides in this PeptideSet
	 * @return number of Peptides
	 */
	public int size(){
		return peptides.size();
	}

	/**
	 * Returns the i-th Peptide with [0,size).
	 * @param i Index of the desired Peptide
	 * @return i-th Peptide
	 */
	public Peptide get(int i){
		return peptides.get(i);
	}

	/**
	 * Returns a Vector with all Peptides in this PeptideSet.
	 * @return Vector of Peptides
	 */
	public Vector<Peptide> getAll() {
		return peptides;
	}

	/**
	 * Reads peptides from a file, either a txt or a mzid. 
	 * 
	 * @param filename complete peptide filename
	 * @throws ExitException
	 */
	public void readPeptides(String filename) throws ExitException{
		if (filename.endsWith(".mzid"))
			readMzIdentML(filename);
		else
			readTxt(filename);
	}

	// parses the peptide txt files
	private void readTxt(String filename) throws ExitException{
		Info info = new Info("read ms peptides");
		int count = 0;

		try {
			BufferedReader pepBuffer = new BufferedReader(new FileReader(new File(filename)));

			String tmpheader = null;
			String line, name = "", desc = "";
			if (pepBuffer.ready()){
				tmpheader = pepBuffer.readLine();
			}
			// check the header
			if (tmpheader.equals(header)){
				while (pepBuffer.ready()){
					count++;
					line = pepBuffer.readLine();					
					try {
						// each line resp peptide is passed to a new Peptide object, which parses the line
						peptides.add(new Peptide(line));
						// some files order peptides by proteins and keep only for the first peptide with same protein the protein information
						// so if there is no protein information, "name" and "desc" hopefully have the once from the prior peptide
						if (peptides.lastElement().getProt_acc().matches("")){
							peptides.lastElement().setProt_acc(name);
							peptides.lastElement().setProt_desc(desc);
						}
						// if there are protein informations they are temporally saved for the next peptides
						else{
							name = peptides.lastElement().getProt_acc();
							desc = peptides.lastElement().getProt_desc();
						}
						// in some prot_desc are gene names resp. gene symbols which can not be extracted during the Peptide object generation
						// just cause of the sometimes missing protein information, so its done afterwards
						peptides.lastElement().extrGS();
					} 
					catch (FormatException e) {
						String message = "unreadable peptide:\t" + line;
						Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", message);
					}

					Configurator.checkBreak();
				}
			}
			else{
				String message = "peptide file doesn't match expected schema!\nexpected header:\t" + header + "\nyour header:\t\t" + tmpheader;
				throw new ExitException(message);
			}

			pepBuffer.close();
		} catch (FileNotFoundException e) {
			throw new ExitException(e.getMessage());
		} catch (IOException e) {
			throw new ExitException(e.getMessage());
		}

		deleteDecoys();
		correctProt_accs();
		info.stop(peptides.size(), count);
	}

	// parses the peptide mzid files
	private void readMzIdentML(String filename) throws ExitException{
		Info info = new Info("read ms peptides");
		int count = 0;

		// the MzIdentML object handles the mzid file and returns a string per peptide 
		// in the same format as in the txt files
		MzIdentML mzreader = new MzIdentML();
		Vector<String> peps = mzreader.load(filename);

		for (String pep : peps){
			try {
				peptides.add(new Peptide(pep));
				peptides.lastElement().extrGS();
				count++;
			} 
			catch (FormatException e) {
				String message = "unreadable peptide:\t" + pep;
				Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", message);
			}

			Configurator.checkBreak();
		}

		deleteDecoys();
		correctProt_accs();
		info.stop(count, peps.size());
	}
	
	// deletes all peptides with decoy proteins (when prot_acc start with "RRR")
	private void deleteDecoys() {
		Iterator<Peptide> it = peptides.iterator();
		int old = peptides.size();
		while (it.hasNext()){
			Peptide p = it.next();
			if (p.getProt_acc().startsWith("RRR")){
				it.remove();
			}
		}
		if (old >= peptides.size()){
			String message = "deleted " + (old - peptides.size()) + " decoy peptides";
			Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", message);
		}
	}
	
	// corrects IPI and GI prot_accs
	private void correctProt_accs(){
		for (Peptide p : peptides){
			String prot_acc = p.getProt_acc();
			// extration of IPIs (e.g. "IPI:IPI00375560.3")
			if (prot_acc.startsWith("IPI:")){
				prot_acc = prot_acc.split("\\.")[0];
				try{
					prot_acc = prot_acc.split(":")[1];
				}
				catch (ArrayIndexOutOfBoundsException e){}
			}
			// extraction of GIs (e.g. "gi|4502027")
			if (prot_acc.startsWith("gi|")){
				prot_acc = prot_acc.substring(3, prot_acc.length());
			}
			p.setProt_acc(prot_acc);
		}
	}

	/**
	 * <p>This method reads useful protein ids from a uniprot id mapping file.</p>
	 * 
	 * <p>The prot_acc (UniProtKB-ID, IPI and GI) is used to find other ids such as UniProtKB-AC, RefSeq and Ensembl_TRS,
	 * which are usually the reference in the gene annotation files (UCSC knownGene and Ensemble ensGene).</p>
	 * 
	 * @param filename uniprod id mapping file
	 * @throws ExitException
	 */
	public void readProteinIds(String filename) throws ExitException{
		Info info = new Info("read protein ids (1)");

		HashMap<String,Vector<Integer>> protmap = getCommonProts("name");

		int lines = 0;
		int count = 0;
		boolean countIt = false;
		String[] splits;

		try {
			BufferedReader rsBuffer = new BufferedReader(new FileReader(new File(filename)));

			while (rsBuffer.ready()){
				splits = rsBuffer.readLine().split("\t");
				countIt = false;
				lines++;
				
				// check the UniprotKB-IDs
				if (protmap.containsKey(splits[1])){
					for (Integer i : protmap.get(splits[1])){
//						if (!peptides.get(i).isProteinAssigned()){
							peptides.get(i).setUniqueIdentifier(splits[0].split("-")[0]);
							for (String ref : splits[3].split("; ")){
								peptides.get(i).addRefSeq(ref.split("\\.")[0]);
							}
							if (splits.length > 20)
								for (String ens : splits[20].split("; ")){
									peptides.get(i).addEnsembl_TRS(ens);
								}
							peptides.get(i).setProteinAssigned(true);
							countIt = true;
//						}
					}
				}
				// check the IPIs, might be more than one per line
				for (String s : splits[7].split("; ")){
					if (protmap.containsKey(s)){
						for (Integer i : protmap.get(s)){
//							if (!peptides.get(i).isProteinAssigned()){
								peptides.get(i).setUniqueIdentifier(splits[0].split("-")[0]);
								for (String ref : splits[3].split("; ")){
									peptides.get(i).addRefSeq(ref.split("\\.")[0]);
								}
								if (splits.length > 20)
									for (String ens : splits[20].split("; ")){
										peptides.get(i).addEnsembl_TRS(ens);
									}
								peptides.get(i).setProteinAssigned(true);
								countIt = true;
//							}
						}
					}
				}
				// check the GIs, might be more than one per line
				for (String s : splits[4].split("; ")){
					if (protmap.containsKey(s)){
						for (Integer i : protmap.get(s)){
//							if (!peptides.get(i).isProteinAssigned()){
								peptides.get(i).setUniqueIdentifier(splits[0].split("-")[0]);
								for (String ref : splits[3].split("; ")){
									peptides.get(i).addRefSeq(ref.split("\\.")[0]);
								}
								if (splits.length > 20)
									for (String ens : splits[20].split("; ")){
										peptides.get(i).addEnsembl_TRS(ens);
									}
								peptides.get(i).setProteinAssigned(true);
								countIt = true;
//							}
						}
					}
				}
				
				if (countIt) count++;
				Configurator.checkBreak();
			}

		} catch (FileNotFoundException e) {
			throw new ExitException(e.getMessage());
		} catch (IOException e) {
			throw new ExitException(e.getMessage());
		}

		info.stop(count, lines);
	}

	/**
	 * <p>This method tries to read some ids from a uniprot fasta file (whole proteome is recommended), particularly from fasta headers.</p>
	 * 
	 * <p>It uses the extracted gene names or gene symbols ({@link Peptide#extrGS()}) to match a potential fasta header and then extracts the UniProtKB-AC.</p>
	 * 
	 * <p>It only involves peptides which have no assigned protein already (e.g. after using {@link #readProteinIds(String)}).</p>
	 * 
	 * @param filename
	 * @throws ExitException
	 */
	public void readFasta(String filename) throws ExitException{

		Vector<Integer> unmapped = new Vector<Integer>();
		for (int p=0; p<size(); p++){
			if (!peptides.get(p).isProteinAssigned()) unmapped.add(p);
		}
		if (unmapped.size() == 0) return;

		Info info = new Info("read protein ids (2)");
		String line;
		int count = 0;
		int peps = 0;

		String id;
		HashMap<String,Vector<Integer>> gnmap = new HashMap<String,Vector<Integer>>();
		for (int i=0; i<unmapped.size(); i++){
			id = peptides.get(unmapped.get(i)).getGN();
			if (!gnmap.containsKey(id))
				gnmap.put(id, new Vector<Integer>());
			gnmap.get(id).add(unmapped.get(i));
		}

		String label = "GN=";
		String gn;
		int pos, from, to;

		try {
			BufferedReader protBuffer = new BufferedReader(new FileReader(new File(filename)));

			while (protBuffer.ready()){
				line = protBuffer.readLine();
				if (line.startsWith(">")){
					peps++;
					pos = line.indexOf(label);
					if (pos >= 0){
						from = pos + label.length();
						to = line.indexOf(" ", from);
						if (to >= 0){
							gn =  line.substring(from, to);
						}
						else{
							gn =  line.substring(from);
						}
						if (gnmap.containsKey(gn)){
							for (Integer i : gnmap.get(gn)){
								if (!peptides.get(i).isProteinAssigned()){
									peptides.get(i).setUniqueIdentifier(line.split("\\|")[1].split("-")[0]);
									peptides.get(i).setProteinAssigned(true);
									count++;
								}
							}
						}
					}
				}	
				Configurator.checkBreak();
			}
			protBuffer.close();

		} catch (FileNotFoundException e) {
			throw new ExitException(e.getMessage());
		} catch (IOException e) {
			throw new ExitException(e.getMessage());
		}

		info.stop(count, peps);
	}

	/**
	 * Providing a HashMap linking peptides to their protein, where a protein identifier is the key and a list of peptides the value.
	 * 
	 * @param ref Parameter to choose the protein identifier, "name" for the protein name (UniProtKB-ID) and "id" for protein ids (UniProtKB-AC and RefSeq).
	 * @return HashMap with proteins as keys and Vectors of peptides as values
	 */
	public HashMap<String, Vector<Integer>> getCommonProts(String ref){
		if (!(ref.matches("name") || ref.matches("id")))
			return null;

		HashMap<String,Vector<Integer>> commonProts = new HashMap<String,Vector<Integer>>();
		Vector<String> id;

		for (int i=0; i<peptides.size(); i++){
			id = peptides.get(i).getProt(ref);
			for (String s : id){
				if (!commonProts.containsKey(s))
					commonProts.put(s, new Vector<Integer>());
				commonProts.get(s).add(i);
			}
		}
		commonProts.remove("");

		return commonProts;
	}

	/**
	 * <p>Writes all mapped Peptides to txt files with same format as the peptide input txt files but extended by the fields
	 * "chrom", "strand", "start_pos", "stop_pos" and "shared". They provide the genome position of the peptides.</p>
	 * 
	 * <p>As a peptide might have several matches, it is exported in several lines, each with the same peptide information but with other positions.</p>
	 * 
	 * <p>There will be two output files regarding the two matching methods (see {@link Mapper}).</p>
	 * 
	 * @param outputPath the path where the files should be written
	 * @param msPepSetName a name for the peptides output, which is used as part of the filename
	 * @throws ExitException
	 */
	public void writeMapped(String outputPath, String msPepSetName) throws ExitException{
		Info info = new Info("write map peptides");
		int count = 0;

		String sep = "\t";

		try {
			BufferedWriter pepBuffer = new BufferedWriter(new FileWriter(new File(outputPath + msPepSetName + "_annomapped.txt")));
			BufferedWriter pepBuffer2 = new BufferedWriter(new FileWriter(new File(outputPath + msPepSetName + "_altmapped.txt")));

			pepBuffer.write(header + sep + "chrom" + sep + "strand" + sep
					+ "start_pos" + sep + "stop_pos" + sep + "shared");
			pepBuffer.newLine();

			pepBuffer2.write(header + sep + "chrom" + sep + "strand" + sep
					+ "start_pos" + sep + "stop_pos" + sep + "shared");
			pepBuffer2.newLine();

			Peptide pep;
			for (int i=0; i<peptides.size(); i++){
				if (peptides.get(i).isGeneMapped()){
					pep = peptides.get(i);
					for (Position pos : pep.getPositions()){
						String line = pep.getRow() + sep
								+ pos.getChrom() + sep + pos.getStrand() + sep
								+ posToString(pos.getStartPos()) + sep + posToString(pos.getEndPos())
								+ sep + pep.getPositions().size();
						if (pep.getMatchMethod().matches(Mapper.ANNOTATION)){
							pepBuffer.write(line);
							pepBuffer.newLine();
						}
						if (pep.getMatchMethod().matches(Mapper.ALTERNATIVE)){
							pepBuffer2.write(line);
							pepBuffer2.newLine();
						}
					}
					count++;
				}
				Configurator.checkBreak();
			}

			pepBuffer.close();
			pepBuffer2.close();
		} catch (IOException e) {
			System.out.println("\nERROR: " + e.getMessage());
		}

		info.stop(count, peptides.size());
	}

	/**
	 * <p>Writes all unmapped Peptides to txt files with same format as the peptide input txt files but extended by the field
	 * "maperror", which tries to indicate whats missing for the mapping.</p>
	 * 
	 * @param outputPath the path where the file should be written
	 * @param msPepSetName a name for the peptides output, which is used as part of the filename
	 * @throws ExitException
	 */
	public void writeUnmapped(String outputPath, String msPepSetName) throws ExitException{
		Info info = new Info("write left peptides");
		int noprot = 0;
		int nogene = 0;
		int noalign = 0;

		String sep = "\t";

		try {
			BufferedWriter pepBuffer = new BufferedWriter(new FileWriter(new File(outputPath + msPepSetName + "_unmapped.txt")));

			pepBuffer.write(header + sep + "maperror");
			pepBuffer.newLine();

			for (int i=0; i<peptides.size(); i++){
				StringBuffer line = new StringBuffer(peptides.get(i).getRow());
				if (!peptides.get(i).isProteinAssigned() && !peptides.get(i).isGeneAssigned() && !peptides.get(i).isGeneMapped()){
					line.append(sep + "noProt");
					pepBuffer.write(line.toString());
					pepBuffer.newLine();
					noprot++;
				}
				if (peptides.get(i).isProteinAssigned() && !peptides.get(i).isGeneAssigned() && !peptides.get(i).isGeneMapped()){
					line.append(sep + "noGene_(prots:" + peptides.get(i).getProt("id") + ")");
					pepBuffer.write(line.toString());
					pepBuffer.newLine();
					nogene++;
				}
				if (peptides.get(i).isProteinAssigned() && peptides.get(i).isGeneAssigned() && !peptides.get(i).isGeneMapped()){
					line.append(sep + "noMatch_(prots:" + peptides.get(i).getProt("id") + ",genes:" + peptides.get(i).getGenes() + ")");
					pepBuffer.write(line.toString());
					pepBuffer.newLine();
					noalign++;
				}
				Configurator.checkBreak();
			}

			pepBuffer.close();
		} catch (IOException e) {
			System.out.println("\nERROR: " + e.getMessage());
		}

		info.stop(noprot+nogene+noalign, peptides.size());
		System.out.println("(no prots: " + noprot + ",no genes: " + nogene + ",no matches: " + noalign + ")");
	}

	// makes a String representation with comma separation of a position vector
	private String posToString(Vector<Integer> pos){
		StringBuffer sb = new StringBuffer();
		for (int i : pos){
			sb.append(i + ",");
		}
		return sb.toString();
	}

	/**
	 * Returns a HashMap indicating which Genes the mapped Peptides have in common.
	 * So the key is a gene name and the value is list of Peptide indices. 
	 * @return a HashMap describing Gene Peptide connections
	 */
	public HashMap<String, Vector<Integer>> getCommonGenes() {
		HashMap<String,Vector<Integer>> commonGenes = new HashMap<String,Vector<Integer>>();
		String id;
		Peptide pep;

		for (int i=0; i<peptides.size(); i++){
			pep = peptides.get(i);
			for (Position pos : pep.getPositions()){
				id = pos.getGeneName();
				if (commonGenes.containsKey(id)){
					commonGenes.get(id).add(i);
				}
				else{
					commonGenes.put(id, new Vector<Integer>());
					commonGenes.get(id).add(i);
				}
			}
		}

		return commonGenes;
	}

	/**
	 * Returns the indices of Peptides which are not mapped yet.
	 * @return Unmapped Peptides indices
	 */
	public Vector<Integer> getUnmapped() {
		Vector<Integer> unmapped = new Vector<Integer>();
		for (int p=0; p<size(); p++){
			if (!peptides.get(p).isGeneMapped()) unmapped.add(p);
		}
		return unmapped;
	}

	/**
	 * <p>Returns a list of indices for a subset of Peptides depending on the "unique" property.</p>
	 * 
	 * <p>Parameters can be either "unique" (returns Peptides with unique value "1"),
	 * "nonunique" (returns Peptides with unique value "0") or "rest" (returns Peptides with an other unique value than "1" or "0").</p>
	 * 
	 * @param param either "unique", "nonunique" or "rest", anything else yields a empty Vector
	 * @return Vector with Peptides indices
	 */
	public Vector<Integer> getUnique(String param) {
		Vector<Integer> unique = new Vector<Integer>();

		if (param.matches("unique")){
			for (int p=0; p<size(); p++){
				if (peptides.get(p).getPep_isunique().matches("1")) unique.add(p);
			}
		}
		if (param.matches("nonunique")){
			for (int p=0; p<size(); p++){
				if (peptides.get(p).getPep_isunique().matches("0")) unique.add(p);
			}
		}
		if (param.matches("rest")){
			for (int p=0; p<size(); p++){
				if (!peptides.get(p).getPep_isunique().matches("1") && !peptides.get(p).getPep_isunique().matches("0"))
					unique.add(p);
			}
		}

		return unique;
	}

	// A Comparator to sort Peptides according to query, sequence and score
	private class PeptideComparator implements Comparator<Peptide> {
		@Override
		public int compare(Peptide pep0, Peptide pep1) {
			Integer query0 = Integer.parseInt(pep0.getPep_query());
			Integer query1 = Integer.parseInt(pep1.getPep_query());
			int result = query0.compareTo(query1);
			if (result == 0){
				String seq0 = pep0.getSequence();
				String seq1 = pep1.getSequence();
				result = seq0.compareTo(seq1);
				if (result == 0){
					Double score0 = Double.parseDouble(pep0.getPep_Score());
					Double score1 = Double.parseDouble(pep1.getPep_Score());
					result = -1 * score0.compareTo(score1);
				}
			}
			return result;
		}
	}

	/**
	 * This method sorts the Peptides to group them according to the pep_query
	 * and gives their Positions a unique String (modifier) each (e.g. "aaa", "aab", "aac").
	 * So in the and each Position will have a unique identifier (pep_query + modifier) for the outputs.
	 */
	public void sortPeptides(){
		Collections.sort(peptides, new PeptideComparator());

		Vector<Peptide> peps = new Vector<Peptide>();
		peps.add(peptides.get(0));
		for (int i=1; i<peptides.size(); i++){
			if (peps.lastElement().getPep_query().matches(peptides.get(i).getPep_query())){
				peps.add(peptides.get(i));
			}
			else{
				modify(peps);
				peps.clear();
				peps.add(peptides.get(i));
			}
		}
		modify(peps);
	}

	// add a unique modifier String to each Position in the passed Peptides
	private void modify(Vector<Peptide> peps){
		int positions = 0;
		for (Peptide pep : peps){
			positions += pep.getPositions().size();
		}

		if (positions>1){
			int digits = (int) Math.ceil(Math.log(positions)/Math.log(26));
			char[] mod = new char[digits];
			Arrays.fill(mod, 'a');
			for (Peptide pep : peps){
				for (Position pos : pep.getPositions()){
					pos.setModifier(new String(mod));

					mod[digits-1]++;
					for (int m=mod.length-1; m>0; m--){
						if (mod[m] == 'a'+26){
							mod[m] = 'a';
							mod[m-1]++;
						}
					}
				}
			}
		}
	}

	/**
	 * Counts the number of Positions (!) of peptides which got mapped to a gene.
	 * 
	 * @return Number of Positions of mapped peptides.
	 */
	public int mapPosCount(){
		int count = 0;	
		for (Peptide pep : peptides){
			count += pep.getPositions().size();
		}
		return count;
	}

	/**
	 * Counts the number of Peptides (!) which got mapped to a gene by a given method.
	 * 
	 * @param method 
	 * @return Number of mapped peptides.
	 */
	public int mapPepCount(String method){
		int count = 0;	
		for (Peptide pep : peptides){
			if (pep.getMatchMethod() != null && pep.getMatchMethod().matches(method))
				count ++;
		}
		return count;
	}
}
