/* 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.File;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map.Entry;
import java.util.Set;
import java.util.TreeMap;
import java.util.Vector;
import java.util.regex.Pattern;

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

/**
 * GeneSet handles the annotation informations provided for UCSC Genes.
 * This includes import methods and creation of a Gene object for each gene imported from a file.
 * 
 * @author Mathias Kuhring
 */
public class GeneSet{

	// The GeneSet's genes
	private Vector<Gene> genes;

	// variables to check the input files
	private final String ucscHeader = "#name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	proteinID	alignID";
	private final String ensembleHeader = "#bin	name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	score	name2	cdsStartStat	cdsEndStat	exonFrames";
	boolean ucsc;
	boolean ensembl;

	/**
	 * <p>Initialization of a GeneSet by giving a filename of an UCSC gene annotations export file
	 * (either with track "UCSC Genes" and table "knownGene" or the track "Ensembl Genes" with table "ensGene").</p>
	 * 
	 * <p>The version with filter {@link #GeneSet(String, Set)} should be preferred as it might improve the overall performance.</p>
	 * 
	 * <p>The file must be tab-separated and start with one of the following headers corresponding to the UCSC table schema:<br>
	 * knownGene:	"#name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	proteinID	alignID"<br>
	 * ensGene:		"#bin	name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	score	name2	cdsStartStat	cdsEndStat	exonFrames"</p>
	 * 
	 * <p>For further informations on the UCSC table schema see {@link Gene} </p>
	 * 
	 * @see #readAnnotations(String filename)
	 * @param filename Name of an UCSC annotation file.
	 * @throws ExitException 
	 */
	public GeneSet(String filename) throws ExitException{
		genes = new Vector<Gene>();
		readAnnotations(filename);
	}

	/**
	 * <p>Initialization of a GeneSet by giving a filename of an UCSC gene annotations export file
	 * (either with track "UCSC Genes" and table "knownGene" or the track "Ensembl Genes" with table "ensGene")
	 * and a Set of protein IDs (UniProtKB-AC and RefSeq) for filtering.</p>
	 * 
	 * <p>The file must be tab-separated and start with one of the following headers corresponding to the UCSC table schema:<br>
	 * knownGene:	"#name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	proteinID	alignID"</p>
	 * ensGene:		"#bin	name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	score	name2	cdsStartStat	cdsEndStat	exonFrames"</p>
	 * 
	 * <p>For further informations on the UCSC table schema see {@link Gene} </p>
	 * 
	 * <p>Only annotations of genes coding for a protein in the filter set will be imported!</p>
	 * 
	 * @see #readAnnotations(String, Set)
	 * @param filename Name of an UCSC annotation file.
	 * @param idFilter Set of Strings, each corresponding to an Uniprot protein ID
	 * @throws ExitException 
	 */
	public GeneSet(String filename, Set<String> idFilter) throws ExitException{
		genes = new Vector<Gene>();
		readAnnotations(filename, idFilter);
	}

	/**
	 * <p>Initialization of a GeneSet by giving a Vector with genes.</p>
	 * 
	 * <p>The genes are added to this new GeneSet, if the all have the same source (either Gene.UCSC or Gene.ENSEMBL).
	 * Otherwise this GeneSet stays empty.</p>
	 * 
	 * @param genes A Vector containing genes ({@link Gene})
	 */
	public GeneSet(Vector<Gene> genes){
		this.genes = new Vector<Gene>();
		if (sameSource(genes)){
			this.genes.addAll(genes);
			ucsc = genes.get(0).getSource().matches(Gene.UCSC);
			ensembl = genes.get(0).getSource().matches(Gene.ENSEMBL);
		}
	}
	
	/**
	 * <p>Adds {@link Gene}s from a Vector to this GeneSet if they have the same source as the present genes.</p>
	 * 
	 * @param input A Vector containing genes ({@link Gene})
	 * @return true if genes were added, else false
	 */
	public boolean addGenes(Vector<Gene> input){
		boolean same;
		if (same = sameSource(input))
			genes.addAll(input);
		return same;
	}

	/*
	 * Checks if the genes in a vector have all the same source among each other 
	 * and as the genes in this GeneSet (only if there are already some).
	 */
	private boolean sameSource(Vector<Gene> input){
		boolean same = true;
		String source;

		if (ucsc) source = Gene.UCSC;
		else if (ensembl) source = Gene.ENSEMBL;
		else source = input.get(0).getSource();

		for (Gene gene : input){
			same &= gene.getSource().matches(source);
		}

		return same;
	}

	/**
	 * <p>Returns the number of Genes in this GeneSet.</p>
	 * <p>Together with the {@link #get(int)}-Method it can be used for loops e.g., cause it just calls the size-method of the vector containing the GeneAnnotations</p>
	 * 
	 * @return Size of the GeneSet, resp. number of Genes.
	 */
	public int size(){
		return genes.size();
	}

	/**
	 * <p>Returns the arg0-th Gene in this GeneSet.</p>
	 * <p>Together with the {@link #size()}-Method it can be used for loops, cause it just calls the get-method of the vector containing the Genes</p>
	 * 
	 * @return arg0-th Gene of the GeneSet.
	 */
	public Gene get(int arg0){
		return genes.get(arg0);
	}

	/**
	 * Returns all genes in this GeneSet in a Vector.
	 * 
	 * @return Vector with {@link Gene}s
	 */
	public Vector<Gene> getAll(){
		return genes;
	}

	/*
	 * Imports a UCSC gene annotation file named with filename.
	 * See also GeneSet(String).
	 */
	private void readAnnotations(String filename) throws ExitException{
		Info info = new Info("read annotations");
		int lineCount = 0;
		int unused = 0;

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

			if (annoBuffer.ready()){
				String tmpheader = annoBuffer.readLine();
				
				// check the header line
				ucsc = tmpheader.matches(ucscHeader);
				ensembl = tmpheader.matches(ensembleHeader);

				// if not a expected header, check if data format fits
				if (!(ucsc || ensembl)){
					ucsc = Gene.isUcscLine(tmpheader.split("\t"));
					ensembl = Gene.isEnsemblLine(tmpheader.split("\t"));
					if (ucsc || ensembl){
						annoBuffer.close();
						annoBuffer = new BufferedReader(new FileReader(new File(filename)));
					}
				}
				
				if (ucsc || ensembl){
					String line;
					while (annoBuffer.ready()){
						line = annoBuffer.readLine();
						
						// check line for correct number of fields
						if ((ucsc && line.split("\t").length != 12) || (ensembl && line.split("\t").length != 16))
							throw new ExitException("error: couldn't parse annotations (wrong # of fields in line " + (lineCount+2) + ")");
						try {
							if (ucsc)
								genes.add(new Gene(line, Gene.UCSC));	
							else if (ensembl)
								genes.add(new Gene(line, Gene.ENSEMBL));	
						} 
						catch (FormatException e) {
							Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", e.getLocalizedMessage());
							unused++;
						}	
						lineCount++;

						Configurator.checkBreak();
					}
				}
				else{
					String message = "Error: annotation file doesn't match UCSC table schema with tab seperation (USCS Genes knownGene or Ensembl Genes ensGene)!";
					throw new ExitException(message);
				}
			}
			else{
				String message = "error: couldn't read file (" + new File(filename).getAbsolutePath() + ")";
				throw new ExitException(message);
			}

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

		info.stop(genes.size(), lineCount);
		if (unused > 0){
			if (Gene.isStrict()){
				System.out.println("(unused genes: " + unused + " (non-protein-coding, uncommon chroms or incorrect/unregocnized annotations))");
			}
			else{
				System.out.println("(unused genes: " + unused + " (incorrect/unregocnized annotations))");
			}
		}
	}

	/*
	 * Imports a UCSC gene annotation file named with filename using a filter set of Uniprot protein IDs.
	 * See also GeneSet(String, Set<String>)
	 */
	private void readAnnotations(String filename, Set<String> idFilter) throws ExitException{
		Info info = new Info("read annotations");
		int lineCount = 0;
		int unused = 0;

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

			if (annoBuffer.ready()){
				String tmpheader = annoBuffer.readLine();

				// check the header line
				ucsc = tmpheader.matches(ucscHeader);
				ensembl = tmpheader.matches(ensembleHeader);

				// if not a expected header, check if data format fits
				if (!(ucsc || ensembl)){
					ucsc = Gene.isUcscLine(tmpheader.split("\t"));
					ensembl = Gene.isEnsemblLine(tmpheader.split("\t"));
					if (ucsc || ensembl){
						annoBuffer.close();
						annoBuffer = new BufferedReader(new FileReader(new File(filename)));
					}
				}
				
				if (ucsc || ensembl){
					String line;
					while (annoBuffer.ready()){
						line = annoBuffer.readLine();

						// check line for correct number of fields
						if ((ucsc && line.split("\t").length != 12) || (ensembl && line.split("\t").length != 16))
							throw new ExitException("error: couldn't parse annotations (wrong # of fields in line " + (lineCount+2) + ")");

						// check if the gene codes for a protein in the filter
						if (ucsc && idFilter.contains(line.split("\t")[10].split("-")[0]) || idFilter.contains(line.split("\t")[1])){ //
							try {
								if (ucsc)
									genes.add(new Gene(line, Gene.UCSC));	
								else if (ensembl)
									genes.add(new Gene(line, Gene.ENSEMBL));	
							} 
							catch (FormatException e) {
								Logger.write(Configurator.getSysProperty("msPepSetName", "ipig")+".log", e.getLocalizedMessage());
								unused++;
							}	
						}
						lineCount++;

						Configurator.checkBreak();
					}
				}
				else{
					String message = "error: annotation file doesn't match UCSC table schema with tab seperation (e.g. USCS Genes knownGene or Ensembl Genes ensGene)!";
					throw new ExitException(message);
				}
			}
			else{
				String message = "error: couldn't read file (" + new File(filename).getAbsolutePath() + ")";
				throw new ExitException(message);
			}

			annoBuffer.close();

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

		info.stop(genes.size(), lineCount);
		if (unused > 0){
			if (Gene.isStrict()){
				System.out.println("(unused genes: " + unused + " (non-protein-coding, uncommon chroms or incorrect/unregocnized annotations))");
			}
			else{
				System.out.println("(unused genes: " + unused + " (incorrect/unregocnized annotations))");
			}
		}
	}

	/**
	 * <p>Prints the annotations of all genes in the set on the screen.</p>
	 * <p>This method is recommended for small sets only, as it would produce a lot of screen output for larger sets.
	 * For large sets use {@link #printAnnotations(int number)} instead to reduce the number of printed annotations.</p>
	 */
	public void printAnnotations(){
		printAnnotations(genes.size());
	}

	/**
	 * <p>Prints the annotations of a chosen number of genes on the screen, starting with the first gene in the set.</p>
	 * 
	 * @param number Number of genes to be printed.
	 */
	public void printAnnotations(int number){
		number = Math.min(number, genes.size());
		System.out.println("GeneAnnotationSet contains " + genes.size() + " gene annotations");
		System.out.println("Printing " + number + " of " + genes.size() + " genes:");
		System.out.println(ucscHeader);
		for (int i=0; i<number; i++){
			System.out.println(genes.get(i).toString());
		}
	}

	/**
	 * <p>Reads the nucleotide sequences for all genes from chromosome fasta files.</p>
	 * 
	 * <p>The chromosome fasta files must be named like the chromosomes in the gene annotations plus the .fa extension
	 * (e.g. "chr1.fa", "chr2.fa", "chrY.fa" and so on, or maybe "chrI.fa", "chrII.fa", "chrIV.fa" etc. depending on the annotation).</p>
	 * 
	 * <p>They must be all located in the same folder, which is indicated as input parameter.</p>
	 * 
	 * <p>The first line in the fasta files must start with ">" followed by the chromosomes name (e.g. ">chr1" or ">chrVI").
	 * The method assumes 50 (!) nucleotides per line, lines with other lengths will cause wrong position calculations and therefore wrong sequences.</p>
	 * 
	 * @param chrompath Folder containing all chromosomes needed for the genes
	 * @throws ExitException 
	 */
	public void readNaSequences(String chrompath) throws ExitException{
		Info info = new Info("read gene sequences");
		int count = 0;

		/* A hashmap assigns a vector of Genes (or actually there indices) to the chromosomes.
		 * The genes in a vector are then sorted by the txStart and txEnd, so running once over the chromosome is enough to read all needed sequences.
		 * Overlaps are also considered. Chromosome parts without needed sequences are skipped.
		 */
		
		HashMap<String,Vector<Integer>> chroms = getCommonChroms();
		HashMap<Integer,StringBuffer> tmpSeqs = new HashMap<Integer,StringBuffer>();
		TreeMap<Integer,Vector<Integer>> txStarts = new TreeMap<Integer,Vector<Integer>>();
		TreeMap<Integer,Vector<Integer>> txEnds = new TreeMap<Integer,Vector<Integer>>();
		Vector<Integer> geneidx;
		String chrom;

		int skip, start, end;
		int skipped = 0;
		int key;
		int idx;

		Entry<Integer,Vector<Integer>> txs, txe;
		String line;

		for (Entry<String, Vector<Integer>> entry : chroms.entrySet()){
			geneidx = entry.getValue();
			chrom = entry.getKey();
			skipped = 0;

			// build sorted lists with txStarts and txEnds with connections to their related genes
			for (int i=0; i<geneidx.size(); i++){
				key = genes.get(geneidx.get(i)).getTxStart();
				if (!txStarts.containsKey(key)) 
					txStarts.put(key, new Vector<Integer>());
				txStarts.get(key).add(geneidx.get(i));

				key = genes.get(geneidx.get(i)).getTxEnd();
				if (!txEnds.containsKey(key)) 
					txEnds.put(key, new Vector<Integer>());
				txEnds.get(key).add(geneidx.get(i));
			}

			try {
				// check path
				if (!chrompath.endsWith("\\") && !chrompath.endsWith("/"))
					chrompath = chrompath + "/";
				BufferedReader chromBuffer = new BufferedReader(new FileReader(new File(chrompath + chrom + ".fa")));

				// check file
				if (chromBuffer.ready()){
					Pattern p = Pattern.compile(">" + chrom);
					if (!chromBuffer.readLine().matches(p.pattern())){	
						continue;	
					}
				}

				while (!txEnds.isEmpty() && chromBuffer.ready()){

					// skip chromosome parts towards next txStart
					if (tmpSeqs.isEmpty()){
						txs = txStarts.firstEntry();
						skip = (int) Math.ceil((txs.getKey()+1) / 50d) - 1 - skipped;
						for (int i=0; i<txs.getValue().size(); i++)
							tmpSeqs.put(txs.getValue().get(i), new StringBuffer());
						chromBuffer.skip(skip*50+skip);
						skipped += skip;
						txStarts.remove(txs.getKey());
					}

					// take next txStarts resp. add gene indices to a sequence buffer, if they are in the next line
					while (!txStarts.isEmpty() && (int) Math.ceil((txStarts.firstKey()+1) / 50d) - 1 == skipped){
						txs = txStarts.firstEntry();
						for (int i=0; i<txs.getValue().size(); i++)
							tmpSeqs.put(txs.getValue().get(i), new StringBuffer());
						txStarts.remove(txs.getKey());
					}

					// read next line and add it to all entries in the sequence buffer
					line = chromBuffer.readLine();
					for (Entry<Integer,StringBuffer> seq : tmpSeqs.entrySet()){
						seq.getValue().append(line);
					}

					// if a gene ends in the current line, finish its sequence with cutting head and tail up to txStart resp. from txEnd
					while (!txEnds.isEmpty() && (int) Math.ceil((txEnds.firstKey()+1) / 50d) - 1 == skipped){
						txe = txEnds.firstEntry();
						for (int i=0; i<txe.getValue().size(); i++){
							idx = txe.getValue().get(i);
							start = genes.get(idx).getTxStart() % 50;
							end = genes.get(idx).getTxEnd() - genes.get(idx).getTxStart() + start;
							genes.get(idx).setSequence(tmpSeqs.remove(idx).toString().substring(start, end));
							count++;
						}
						txEnds.remove(txe.getKey());
					}

					skipped++;

					Configurator.checkBreak();
				}

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

		}

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

	/**
	 * Deletes the nucleotide sequence of every gene.
	 */
	public void deleteNaSequences(){
		for (Gene gene : genes){
			gene.setSequence(null);
		}
		System.gc();
	}

	/**
	 * <p>This function calculates the proteins the genes have in common.</p>
	 * 
	 * <p>It will return a HashMap with protein identifiers as keys and as value a Vector 
	 * containing the indices of genes (in this GeneSet) which code for this protein.</p>
	 *  
	 * @return a HashMap with proteins as keys and Vectors of gene indices as values
	 */
	public HashMap<String, Vector<Integer>> getCommonProts(){
		HashMap<String,Vector<Integer>> commonProts = new HashMap<String,Vector<Integer>>();
		String id;

		for (int i=0; i<genes.size(); i++){
			id = genes.get(i).getProteinReference();
			if (!commonProts.containsKey(id)){
				commonProts.put(id, new Vector<Integer>());
			}
			commonProts.get(id).add(i);
		}

		return commonProts;
	}

	/**
	 * <p>This function calculates the chromosomes the genes have in common.</p>
	 * 
	 * <p>It will return a HashMap with chromosome names as keys and as value a Vector 
	 * containing the indices of genes (in this GeneSet) which lie on this chromosome.</p>
	 *  
	 * @return a HashMap with chromosomes as keys and Vectors of gene indices as values
	 */
	public HashMap<String, Vector<Integer>> getCommonChroms(){
		HashMap<String,Vector<Integer>> commonChroms = new HashMap<String,Vector<Integer>>();
		String id;

		for (int i=0; i<genes.size(); i++){
			id = genes.get(i).getChrom();
			if (!commonChroms.containsKey(id)){
				commonChroms.put(id, new Vector<Integer>());
			}
			commonChroms.get(id).add(i);
		}

		return commonChroms;
	}

	/**
	 * <p>This functions reads the aminoacid sequences related to the genes as they can be obtained
	 * from the UCSC table browser (either with track "UCSC Genes" and table "knownGenePep" or the track "Ensembl Genes" with table "ensPep").</p>
	 * 
	 * <p> The header of a aa sequences file has to be "#name	seq"</p>
	 * 
	 * @param filename
	 * @throws ExitException 
	 */
	public void readAaSequences(String filename) throws ExitException{
		Info info = new Info("read aa sequences");
		int count = 0;
		int lines = 0;
//		String tmpheader = null;
		String[] line;

		// create a map with <gene name, gene index>
		HashMap<String,Integer> genemap = new HashMap<String,Integer>();
		for (int g=0; g<genes.size(); g++){
			genemap.put(genes.get(g).getName(),g);
		}

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

//			if (aaBuffer.ready()){
				//  !!! header test canceled, because ucsc files from ftp server are without header !!!
				//	tmpheader = aaBuffer.readLine();
				//	if (tmpheader.matches("#name	seq")){
					while (aaBuffer.ready()){
						lines++;
						line = aaBuffer.readLine().split("\t");
						
						// if the genemap contains the gene name in this line, then the aa sequences is added to the corresponding gene
						if (genemap.containsKey(line[0])){
							genes.get(genemap.get(line[0])).setAaSequence(line[1]);
							count++;
						}

						Configurator.checkBreak();
					}
//				}
//				else{
//					String message = "error: file doesn't match UCSC/Ensembl pep table schema with tab seperation!"
//							+ "\nexpected header:\t" + "#name	seq" 
//							+ "\nyour header:\t\t" + tmpheader;
//					throw new ExitException(message);
//				}
//			}

			aaBuffer.close();

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

		info.stop(count, lines);
		removeAaSeqless();
	}

	/**
	 * Deletes the aminoacid sequence of every gene.
	 */
	public void deleteAaSequences(){
		for (Gene gene : genes){
			gene.setAaSequence(null);
		}
		System.gc();
	}

	/* removes all genes without a aminoacid sequence
	 * e.g. used in readAaSequences after loading the sequences, because genes without are useless in the mapping steps
	 */
	private void removeAaSeqless(){
		Gene gene;
		for (Iterator<Gene> it = genes.iterator(); it.hasNext();){
			gene = it.next();
			if (gene.getAaSequence() == null){
				it.remove();
			}
		}
		System.gc();
	}

	/**
	 * <p>Deletes every Gene which is not marked as used (see {@link Gene#setUsed(boolean)}).</p>
	 * 
	 * <p>This function might be used in front of memory expensive operations like {@link #readNaSequences(String)}.</p>
	 */
	public void removeUnused(){
		for (Iterator<Gene> it = genes.iterator(); it.hasNext();){
			if (!it.next().isUsed()){
				it.remove();
			}
		}
		System.gc();
	}

	/**
	 * Gives a vector with randomly chosen Gene objects.
	 * 
	 * @param number The number of required Gene objects.
	 * @return Vector of Gene objects.
	 */
	public Vector<Gene> getRandomSubset(int number){
		Collections.shuffle(genes);	
		return new Vector<Gene>(genes.subList(0, Math.min(number, genes.size())));
	}	

	public Vector<Gene> getSubset(int start, int end){
		Vector<Gene> subset = new Vector<Gene>();

		end = Math.min(end, genes.size());
		for(int i=start; i<end; i++){
			subset.add(genes.get(i));
		}

		return subset;
	}	

	/**
	 * Returns a String with the header information like in the input file (either ucsc or ensembl genes).
	 * 
	 * @return  the header as a String
	 */
	public String getHeader(){
		if (ucsc)
			return ucscHeader;
		else // ensembl
			return ensembleHeader;
	}
}
