/* 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.FileWriter;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map.Entry;
import java.util.Vector;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

import org.jdom.Document;
import org.jdom.Element;
import org.jdom.JDOMException;
import org.jdom.Namespace;
import org.jdom.input.SAXBuilder;

public class MzIdentML {
	
	private final String xsd1 = "mzIdentML1.0.0.xsd";
	private final String xsd2 = "FuGElightv1.0.0.xsd";
	private final String xsdPath = "/xsd/";
	
	/**
	 * <p>Reads an mzIdentML file and parses the peptide identifications.</p>
	 * 
	 * <p>The output is a Vector with Strings each representing a identified peptide in the same format as expected from the txt file:<br>
	 * "prot_acc	prot_desc	pep_query	pep_isunique	pep_exp_z	pep_score	pep_seq	pep_var_mod	pep_var_mod_pos"</p>
	 * 
	 * @param filename Path of the mzIdentML file
	 * @return a Vector with a String for each peptide identification
	 */
	public Vector<String> load(String filename){
		Vector<String> peptides = null;
				
		try {
			// extract xsds from jar-resources
			xsdExtract();
			
			// load and validate the mzIdentML
			SAXBuilder builder =
					new SAXBuilder("org.apache.xerces.parsers.SAXParser", true);
			builder.setFeature(
					"http://apache.org/xml/features/validation/schema", true);
			builder.setProperty(
					"http://apache.org/xml/properties/schema/external-schemaLocation",
					"http://psidev.info/psi/pi/mzIdentML/1.0 " + new File("mzIdentML1.0.0.xsd").getAbsolutePath());
			Document doc = builder.build(filename);
			
			// delete the xsds
			xsdDelete();

			// start navigating trough xml structure
			Element root = doc.getRootElement();
			Namespace ns = root.getNamespace();

			// get proteins and peptides
			Element seqs = root.getChild("SequenceCollection", ns);
			HashMap<String,Protein> prots = getProts(seqs, ns);
			HashMap<String,Peptide> peps = getPeps(seqs, ns);

			
			// get results
			Element idents = root.getChild("DataCollection", ns).getChild("AnalysisData", ns).getChild("SpectrumIdentificationList", ns);
			Vector<Result> res = getResults(idents, ns);

			// combine them
			peptides = combine(res, peps, prots);

			// write them to a file (just for verification)
//			BufferedWriter pepBuffer = new BufferedWriter(new FileWriter(new File("mzidParsingRes.txt")));
//			pepBuffer.write("prot_acc	prot_desc	pep_query	pep_isunique	pep_exp_z	pep_score	pep_seq	pep_var_mod	pep_var_mod_pos");
//			pepBuffer.newLine();
//			for (String p : peptides){
//				pepBuffer.write(p);
//				pepBuffer.newLine();
//			}
//			pepBuffer.close();

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

		return peptides;
	}
	
	// start of extraction of xsd file from jar-resources, 
	// cause the SAXBuilder need a full usable path as String for xml-Validation (So can't use a resource as stream).
	private String xsdExtract(){
		File f1 = new File(xsd1);
		File f2 = new File(xsd2);
		
		xsdCopy(xsdPath + xsd1, f1);
		xsdCopy(xsdPath + xsd2, f2);
		
		return f1.getAbsolutePath();
	}
	
	// xsd resource is copied via a resource stream to an external file.
	private void xsdCopy(String xsdRes, File xsdFile){
		InputStream is = getClass().getResourceAsStream(xsdRes);

		BufferedReader br = new BufferedReader(new InputStreamReader(is));
		try {
			BufferedWriter bw = new BufferedWriter(new FileWriter(xsdFile));

			while (br.ready()){
				bw.write(br.readLine());
				bw.newLine();
			}

			br.close();
			bw.close();
		} 
		catch (IOException e) {
		}
	}
	
	// external file can be deleted after xml validation
	private void xsdDelete(){
		new File(xsd1).delete();
		new File(xsd2).delete();
	}

	// combines the results with the peptides and proteins to a string per peptide like needed as input for a Peptide-Object
	// prot_acc	prot_desc	pep_query	pep_isunique	pep_exp_z	pep_score	pep_seq	pep_var_mod	pep_var_mod_pos
	private Vector<String> combine(Vector<Result> res, HashMap<String,Peptide> peps, HashMap<String,Protein> prots){
		Vector<String> peptideRows = new Vector<String>();
		String row;
		String sep = "\t";

		for (Result r : res){
			Peptide pep = peps.get(r.Peptide_ref);
			for (String evi : r.DBSequence_Ref){
				row = prots.get(evi).accession + sep + prots.get(evi).proteinDescription + sep + query(r.query) + sep 
						+ unique(r.unique) + sep + r.chargeState + sep + r.score + sep + pep.sequence + sep
						+ mod(pep.sequence, pep.modLocation, pep.modName); 
				peptideRows.add(row);
			}
		}

		return peptideRows;
	}

	// parses the last number in query (originally SpectrumIdentificationResult's "id")
	private String query(String query){
		Matcher matcher = Pattern.compile( "\\d+" ).matcher( query ); 
		while ( matcher.find() ) 
			query = matcher.group();
		return query;
	}

	// returns the string representation for the unique value
	private String unique(boolean unique){
		if (unique) return "1";
		else return "0";
	}

	// builds a string representation of a peptide's modifications similar to mascot csv exports
	private String mod(String seq, Vector<String> locations, Vector<String> names){
		char[] posArray = new char[seq.length()+2];
		Arrays.fill(posArray, '0'); 
		StringBuffer posString = new StringBuffer().append(posArray);
		HashMap<String,Integer> modCount = new HashMap<String,Integer>();

		boolean mod = false;
		int shift = 0;
		for (int i=0; i<locations.size(); i++){
			try{
				int pos = Integer.parseInt(locations.get(i)) + shift;
				
				String num = "X";
				Matcher matcher = Pattern.compile( "UNIMOD:\\d+" ).matcher( names.get(i) );
				while ( matcher.find() ) 
					num = matcher.group().split(":")[1];
				
				posString.replace(pos, pos+1, "[" + num + "]");
				shift += num.length()+1; // == (num.length-1) for nums > 9, +2 for the brackets
				mod = true;
				if (!modCount.containsKey(names.get(i))) 
					modCount.put(names.get(i), 0);
				modCount.put(names.get(i), modCount.get(names.get(i))+1);
			}
			catch (NumberFormatException e){}
		}
		posString.insert(1, ".");
		posString.insert(posString.length()-1, ".");

		StringBuffer nameString = new StringBuffer();
		for (Entry<String,Integer> e : modCount.entrySet()){
			if (!e.getKey().matches(""))
				if (e.getValue()>1)
					nameString.append(e.getValue() + " " + e.getKey() + ", ");
				else
					nameString.append(e.getKey() + ", ");
		}

		if (mod)
			return nameString.substring(0, nameString.length()-2) + "\t" + posString.toString();
		else
			return "\t";
	}

	// parses all proteins
	private HashMap<String,Protein> getProts(Element seqs, Namespace ns){
		@SuppressWarnings("unchecked")
		List<Element> prots = seqs.getChildren("DBSequence", ns);
		HashMap<String,Protein> proteins = new HashMap<String,Protein>();
		for (Element prot : prots){
			String id = prot.getAttributeValue("id");
			Protein temp = new Protein();
			temp.accession = prot.getAttributeValue("accession");
			temp.proteinDescription = "";
			@SuppressWarnings("unchecked")
			List<Element> cvParams = prot.getChildren("cvParam", ns);
			for (Element cv : cvParams){
				if (cv.getAttributeValue("name").matches("protein description"))
					temp.proteinDescription = cv.getAttributeValue("value");
			}
			proteins.put(id, temp);
		}
		return proteins;
	}

	// parses all peptides
	private HashMap<String,Peptide> getPeps(Element seqs, Namespace ns){
		@SuppressWarnings("unchecked")
		List<Element> peps = seqs.getChildren("Peptide", ns);
		HashMap<String,Peptide> peptides = new HashMap<String,Peptide>();
		for (Element pep : peps){
			String id = pep.getAttributeValue("id");
			Peptide temp = new Peptide();
			temp.sequence = pep.getChild("peptideSequence", ns).getValue();
			temp.modLocation = new Vector<String>();
			temp.modName = new Vector<String>();
			@SuppressWarnings("unchecked")
			List<Element> mods = pep.getChildren("Modification", ns);
			for (Element mod : mods){
				temp.modLocation.add(mod.getAttributeValue("location"));
				temp.modName.add("");
				@SuppressWarnings("unchecked")
				List<Element> cvParams = mod.getChildren("cvParam", ns);
				for (Element cv : cvParams){
					if (cv.getAttributeValue("accession").startsWith("UNIMOD")){
						temp.modName.set(temp.modName.size()-1, cv.getAttributeValue("name") + " (" + cv.getAttributeValue("accession") + ")");
					}
				}
			}
			peptides.put(id, temp);
		}
		return peptides;
	}

	// parses the results resp. each SpectrumIdentificationItem containing at least one PeptideEvidence
	private Vector<Result> getResults(Element idents, Namespace ns){
		@SuppressWarnings("unchecked")
		List<Element> res = idents.getChildren("SpectrumIdentificationResult", ns);
		Vector<Result> results = new Vector<Result>();
		for (Element r : res){
			@SuppressWarnings("unchecked")
			List<Element> items = r.getChildren("SpectrumIdentificationItem", ns);
			for (Element item : items){
				@SuppressWarnings("unchecked")
				List<Element> evis = item.getChildren("PeptideEvidence", ns);
				if (evis.size() > 0){
					Result temp = new Result();
					temp.query = r.getAttributeValue("id");
					temp.chargeState = item.getAttributeValue("chargeState");
					temp.Peptide_ref = item.getAttributeValue("Peptide_ref");
					temp.DBSequence_Ref = new Vector<String>();
					for (Element evi : evis){
						temp.DBSequence_Ref.add(evi.getAttributeValue("DBSequence_Ref"));
					}
					temp.score = "0";
					temp.unique = true;
					@SuppressWarnings("unchecked")
					List<Element> cvParams = item.getChildren("cvParam", ns);
					for (Element para : cvParams){
						if (para.getAttributeValue("name").matches("mascot:score"))
							temp.score = para.getAttributeValue("value");
						if (para.getAttributeValue("name").matches("peptide shared in multiple proteins"))
							temp.unique = false;
					}
					results.add(temp);
				}
			}
		}
		return results;
	}

	// class for collecting the protein informations temporally
	private class Protein{
		String accession;
		String proteinDescription;
	}

	// class for collecting the peptide informations temporally
	private class Peptide{
		String sequence;
		Vector<String> modLocation;
		Vector<String> modName;

	}
	
	// class for collecting the result informations temporally
	private class Result{
		String query;
		String chargeState;
		String Peptide_ref;
		Vector<String> DBSequence_Ref;
		boolean unique;
		String score;
	}
}
