/* 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 de.rki.ng4.ipig.exceptions.FormatException;

/**
 * <p>Gene is a container for diverse informations of a gene, as they are provided by the USCS table browser.</p>
 * 
 * <p>The import of USCS annotations should be done with a {@link GeneSet}, which provides reading methods 
 * and creates an Gene object for any gene (resp. line) in an imported annotation file.</p>
 * 
 * <p>Please note: Coordinates (resp. positions) provided by the UCSC are zero-based and half-open.
 * This means gene positions are indexed like for example Java Arrays from 0 to length-1.</p>
 * 
 * <p>The member variables correspond to the USCS table schema for the track "UCSC Genes" and table "knownGene"
 * or the track "Ensembl Genes" with table "ensGene"
 * with e.g. clade "Mammal", genome "Human" and group "Genes and Gene Prediction Tracks" (
 * <a href="http://genome.ucsc.edu/cgi-bin/hgTables?org=Human&db=hg19&hgsid=205072007&hgta_doMainPage=1">
 * USCS Table Browser</a>).</p>
 *
 * @author Mathias Kuhring
 */

/* 
 * UCSC Schema for UCSC Genes
 * field		example				SQL type 			description
 * name 		uc001aaa.3			varchar(255) 		Name of gene
 * chrom 		chr1				varchar(255) 		Reference sequence chromosome or scaffold
 * strand 		+					char(1) 			+ or - for strand
 * txStart 		11873				int(10) unsigned 	Transcription start position
 * txEnd 		14409				int(10) unsigned 	Transcription end position
 * cdsStart 	11873				int(10) unsigned 	Coding region start
 * cdsEnd 		11873				int(10) unsigned 	Coding region end
 * exonCount 	3					int(10) unsigned 	Number of exons
 * exonStarts 	11873,12612,13220,	longblob 	  		Exon start positions
 * exonEnds 	12227,12721,14409,	longblob 	  		Exon end positions
 * proteinID  						varchar(40) 		UniProt display ID for Known Genes, UniProt accession or RefSeq protein ID for UCSC Genes
 * alignID 		uc001aaa.3			varchar(255) 		Unique identifier for each (known gene, alignment position) pair
 * 
 * UCSC Schema for Ensembl Genes
 * field			example			SQL type 				info 	description
 * bin 				585				smallint(5) unsigned 	range 	Indexing field to speed chromosome range queries.
 * name 			ENST00000456328	varchar(255) 			values 	Name of gene
 * chrom 			chr1			varchar(255) 			values 	Reference sequence chromosome or scaffold
 * strand 			+				char(1) 				values 	+ or - for strand
 * txStart 			11868			int(10) unsigned 		range 	Transcription start position
 * txEnd 			14409			int(10) unsigned 		range 	Transcription end position
 * cdsStart 		14409			int(10) unsigned 		range 	Coding region start
 * cdsEnd 			14409			int(10) unsigned 		range 	Coding region end
 * exonCount 		3				int(10) unsigned 		range 	Number of exons
 * exonStarts 		11868,12612,13220,	longblob 	  				Exon start positions
 * exonEnds 		12227,12721,14409,	longblob 	  				Exon end positions
 * score 			0				int(11) 				range 	 
 * name2 			ENSG00000223972	varchar(255) 			values 	 
 * cdsStartStat 	none			enum('none', 'unk', 'incmpl', 'cmpl') 	values 	 
 * cdsEndStat 		none			enum('none', 'unk', 'incmpl', 'cmpl') 	values 	 
 * exonFrames 		-1,-1,-1,		longblob 	  	 
 */
public class Gene{

	// gene source marker
	public static final String UCSC = "ucsc";
	public static final String ENSEMBL = "ensembl";
	
	// if strict == true, imported genes are filtered for non-codings and uncommon proteins
	// strict can be set with setStrict(), e.g. useful for GeneControl, so it can check this stuff itself
	private static boolean strict = true;

	// marks the gene source, either UCSC or ENSEMBLE
	private String source;

	// gene annotation (ucsc & ensemble)
	private String name;
	private String chrom;
	private char strand;
	private int txStart;
	private int txEnd;
	private int cdsStart;
	private int cdsEnd;
	private int exonCount;
	private int exonStarts[];
	private int exonEnds[];

	// gene annotation (ucsc)
	private String proteinID;
	private String alignID;

	// gene annotation (ensemble)
	private int bin;
	private int score;
	private String name2;
	private String cdsStartStat;
	private String cdsEndStat;
	private String exonFrames;

	// sequences
	private String sequence;
	private String aaSequence;

	// marks if a gene was used, that is if any peptide was mapped to this gene
	private boolean used;

	/**
	 * Gene is initialized by processing of an UCSC annotation file row to set the member variables
	 * 
	 * @param row String containing one gene's annotations, so one line of the UCSC annotation file.
	 * @throws FormatException 
	 */
	public Gene(String row, String source) throws FormatException{
		this.source = source;

		String[] splits = row.split("\t");
		checkRow(splits);
		parseRow(splits);
	}

	/*
	 * Checks if the row fits the expected format, either an UCSC or Ensembl gene.
	 * Genes on uncommon chroms (e.g. chr6_qbl_hap6) and non-protein-coding genes are rejected.
	 * All fieds are tested with reg. Expressions.
	 */
	private void checkRow(String[] splits) throws FormatException{
		if (source.matches(UCSC)){
			/* Example:
			 * #name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	proteinID	alignID
			 * uc001aaa.3	chr1	+	11873	14409	11873	11873	3	11873,12612,13220,	12227,12721,14409,		uc001aaa.3
			 * uc010nxq.1	chr1	+	11873	14409	12189	13639	3	11873,12594,13402,	12227,12721,14409,	B7ZGX9	uc010nxq.1
			 * uc001aut.1	chr1	-	13328195	13331692	13328832	13331671	3	13328195,13330413,13331378,	13329406,13330992,13331692,	NP_001094101	uc001aut.1
			 * uc011ihh.1	chr6_qbl_hap6	-	1315373	1321962	1317738	1318302	3	1315373,1318888,1321856,	1318491,1318964,1321962,	Q2KJ03	uc011ihh.1
			 */
			if (strict && !splits[1].matches("chr(\\d+|X|Y|M|[IVX]+)")){
				throw new FormatException("uncommon chromosome: " + splits[0] + " " + splits[1]);
			}
			if (strict && splits[5].matches(splits[6])){
				throw new FormatException("non-protein-coding gene: " + splits[0]);	
			}
			if (!isUcscLine(splits)){
				throw new FormatException("incorrect annotation row: " + splits[0]);
			}
			
		}
		else if (source.matches(ENSEMBL)){
			/* Example
			 * #bin	name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	score	name2	cdsStartStat	cdsEndStat	exonFrames
			 * 0	ENST00000237247	chr1	+	66999065	67210057	67000041	67208778	27	66999065,66999928,67091529,67098752,67099762,67105459,67108492,67109226,67126195,67133212,67136677,67137626,67138963,67142686,67145360,67147551,67149789,67154830,67155872,67161116,67184976,67194946,67199430,67205017,67206340,67206954,67208755,	66999090,67000051,67091593,67098777,67099846,67105516,67108547,67109402,67126207,67133224,67136702,67137678,67139049,67142779,67145435,67148052,67149870,67154958,67155999,67161176,67185088,67195102,67199563,67205220,67206405,67207119,67210057,	0	ENSG00000118473	cmpl	cmpl	-1,0,1,2,0,0,0,1,0,0,0,1,2,1,1,1,1,1,0,1,1,2,2,0,2,1,1,
			 * 0	ENST00000371039	chr1	+	66999274	67210768	67000041	67208778	22	66999274,66999928,67091529,67098752,67105459,67108492,67109226,67136677,67137626,67138963,67142686,67145360,67154830,67155872,67160121,67184976,67194946,67199430,67205017,67206340,67206954,67208755,	66999355,67000051,67091593,67098777,67105516,67108547,67109402,67136702,67137678,67139049,67142779,67145435,67154958,67155999,67160187,67185088,67195102,67199563,67205220,67206405,67207119,67210768,	0	ENSG00000118473	cmpl	cmpl	-1,0,1,2,0,0,1,0,1,2,1,1,1,0,1,1,2,2,0,2,1,1,
			 * 0	ENST00000424320	chr1	+	66999297	67145425	67000041	67145425	13	66999297,66999928,67091529,67098752,67101626,67105459,67108492,67109226,67136677,67137626,67138963,67142686,67145360,	66999355,67000051,67091593,67098777,67101698,67105516,67108547,67109402,67136702,67137678,67139049,67142779,67145425,	0	ENSG00000118473	cmpl	incmpl	-1,0,1,2,0,0,0,1,0,1,2,1,1,
			 * 
			 * #bin	name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	score	name2	cdsStartStat	cdsEndStat	exonFrames
			 * 73	YAL012W	chrI	+	130798	131983	130798	131983	1	130798,	131983,	0	YAL012W	cmpl	cmpl	0,
			 */
			if (strict && !splits[2].matches("chr(\\d+|X|Y|M|[IVX]+)")){
				throw new FormatException("unusable chromosome: " + splits[1] + " " + splits[2]);
			}
			if (strict && splits[6].matches(splits[7])){
				throw new FormatException("non-protein-coding gene: " + splits[1]);	
			}
			if (!isEnsemblLine(splits)){
				throw new FormatException("incorrect annotation row: " + splits[1]);
			}
			
		}
		else
			throw new FormatException("unusable source: " + source);
	}
	
	public static boolean isUcscLine(String[] splits){
		return splits.length == 12 &&
				splits[0].matches("\\w+\\.\\d") && splits[1].matches("\\w+") &&
				splits[2].matches("[\\+\\-]") && splits[3].matches("\\d+") &&
				splits[4].matches("\\d+") && splits[5].matches("\\d+") &&
				splits[6].matches("\\d+") && splits[7].matches("\\d+") &&
				splits[8].matches("(\\d+,)++") && splits[9].matches("(\\d+,)++") &&
				splits[10].matches("[\\w\\-]*") && splits[11].matches("\\w+\\.\\d");
	}
	
	public static boolean isEnsemblLine(String[] splits){
		return splits.length == 16 &&
				splits[0].matches("\\d+") && splits[1].matches("[\\w\\-\\(\\)]+") && splits[2].matches("\\w+") &&
				splits[3].matches("[\\+\\-]") && splits[4].matches("\\d+") &&
				splits[5].matches("\\d+") && splits[6].matches("\\d+") &&
				splits[7].matches("\\d+") && splits[8].matches("\\d+") &&
				splits[9].matches("(\\d+,)++") && splits[10].matches("(\\d+,)++") &&
				splits[11].matches("\\d+") && splits[12].matches("[\\w\\-\\(\\)]+") &&
				splits[13].matches("\\w+") && splits[14].matches("\\w+") &&
				splits[15].matches("(-?\\d+,)+");
	}

	/*
	 * Parses the annotation fields and sets the corresponding variables
	 */
	private void parseRow(String[] splits){	
		if (source.matches(UCSC)){
			// #name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	proteinID	alignID
			int index = 0;
			
			name = splits[index++];
			chrom = splits[index++];
			strand = splits[index++].charAt(0);
			txStart = Integer.parseInt(splits[index++]);
			txEnd = Integer.parseInt(splits[index++]);
			cdsStart = Integer.parseInt(splits[index++]);
			cdsEnd = Integer.parseInt(splits[index++]);
			exonCount = Integer.parseInt(splits[index++]);

			exonStarts = new int[exonCount];
			exonEnds = new int[exonCount];
			String starts[] = splits[index++].split(",");
			String ends[] = splits[index++].split(",");
			for (int i=0; i<exonCount; i++){
				exonStarts[i] = Integer.parseInt(starts[i]);
				exonEnds[i] = Integer.parseInt(ends[i]);
			}

			// import protein id without isoform number (in case of UniProtKB-AC) -> split("-")
			proteinID = splits[index++].split("-")[0];
			alignID = splits[index++];
		}
		
		if (source.matches(ENSEMBL)){
			// #bin	name	chrom	strand	txStart	txEnd	cdsStart	cdsEnd	exonCount	exonStarts	exonEnds	score	name2	cdsStartStat	cdsEndStat	exonFrames
			int index = 0;
			
			bin = Integer.parseInt(splits[index++]);
			name = splits[index++];
			chrom = splits[index++];
			strand = splits[index++].charAt(0);
			txStart = Integer.parseInt(splits[index++]);
			txEnd = Integer.parseInt(splits[index++]);
			cdsStart = Integer.parseInt(splits[index++]);
			cdsEnd = Integer.parseInt(splits[index++]);
			exonCount = Integer.parseInt(splits[index++]);

			exonStarts = new int[exonCount];
			exonEnds = new int[exonCount];
			String starts[] = splits[index++].split(",");
			String ends[] = splits[index++].split(",");
			for (int i=0; i<exonCount; i++){
				exonStarts[i] = Integer.parseInt(starts[i]);
				exonEnds[i] = Integer.parseInt(ends[i]);
			}
			
			score = Integer.parseInt(splits[index++]);
			name2 = splits[index++];
			cdsStartStat = splits[index++];
			cdsEndStat = splits[index++];
			exonFrames = splits[index++];
		}
	}

	/**
	 * Returns a String aggregating all annotations in one line (row) in the same format as the gene had in the UCSC annotation file.
	 * 
	 * @see java.lang.Object#toString()
	 * @return Annotations in a one line String
	 */
	public String toString(){
		if (source.matches(UCSC))
			return name + "\t" + chrom + "\t" + strand + "\t" + txStart + "\t" + txEnd
				+ "\t" + cdsStart + "\t" + cdsEnd + "\t" + exonCount + "\t" + exons() + "\t" + proteinID + "\t" + alignID;
		else // source.matches(ENSEMBL)
			return bin + "\t" + name + "\t" + chrom + "\t" + strand + "\t" + txStart + "\t" + txEnd
					+ "\t" + cdsStart + "\t" + cdsEnd + "\t" + exonCount + "\t" + exons() + "\t" + score
					+ "\t" + name2 + "\t" + cdsStartStat + "\t" + cdsEndStat + "\t" +  exonFrames;
	}

	// returns the exon start and end list as string for gene export (e.g. toString())
	private String exons(){
		StringBuffer starts = new StringBuffer();
		StringBuffer ends = new StringBuffer();

		for (int i=0; i<exonStarts.length; i++){
			starts.append(exonStarts[i] + ",");
			ends.append(exonEnds[i] + ",");
		}

		return starts.toString() + "\t" + ends.toString();
	}

	/**
	 * Returns the name of the gene
	 * 
	 * @return Name of the gene
	 */
	public String getName(){
		return name;
	}

	/**
	 * Returns the chromosome the gene is located on as String (e.g. "chr6").
	 * 
	 * @return String describing a chromosome
	 */
	public String getChrom(){
		return chrom;
	}

	/**
	 * Returns the dna strand the gene is located on, resp. the orientation of the gene on the chromosome.
	 * 
	 * Annotation data is usually provided only for resp. as one single strand in 5'->3' direction.
	 * 
	 * @return '+' for the 5'->3' or '-' for 3'->5' direction
	 */
	public char getStrand(){
		return strand;
	}

	/**
	 * Returns the transcription start position.
	 * 
	 * @return Transcription start position.
	 */
	public int getTxStart(){
		return txStart;
	}

	/**
	 * Returns the transcription end position.
	 * 
	 * @return Transcription end position.
	 */
	public int getTxEnd(){
		return txEnd;
	}

	/**
	 * Returns the start position of the coding sequence, so usually position of the Start-Codon if strand '+', resp. Stop-Codon if strand '-'.
	 * 
	 * @return Position of CdsStart
	 */
	public int getCdsStart(){
		return cdsStart;
	}

	/**
	 * Returns the end position of the coding sequence, so usually position after the stop-codon if strand '+', resp. start-codon if strand '-'.
	 * (Remember: zero-based and half-open, see {@link Gene}).
	 * 
	 * @return Position of CdsEnd
	 */
	public int getCdsEnd(){
		return cdsEnd;
	}

	/**
	 * Returns the number of exons.
	 * 
	 * @return Number of exons
	 */
	public int getExonCount(){
		return exonCount;
	}

	/**
	 * Returns all exon start positions.
	 * 
	 * @return List of exon start positions
	 */
	public int[] getExonStarts(){
		return exonStarts;
	}

	/**
	 * Returns all exon end positions.
	 * (Remember: zero-based and half-open, see {@link Gene}).
	 * 
	 * @return List of exon end positions
	 */
	public int[] getExonEnds(){
		return exonEnds;
	}

	/**
	 * Returns the protein reference (UniProtKB-AC or RefSeq for UCSC, name (corr. Ensembl_TRS) for Ensembl) of the protein the gene encodes.
	 * @return Protein reference string
	 */
	public String getProteinReference(){
		if (source.matches(UCSC))
			return proteinID;
		else // source.matches(ENSEMBL)
			return name;
	}

	/**
	 * <p>Sets the gene's DNA sequence, which must range from the transcription start position (including) to the transcription end position (excluding).</p>
	 * 
	 * <p>Should be used carefully because later calculations may depend on the sequence's correctness, especially regarding the correct relation to txStart and txEnd.</p>
	 * 
	 * @param sequence The gene's DNA sequence as String
	 */
	public void setSequence(String sequence) {
		this.sequence = sequence;
	}

	/**
	 * Returns the gene's DNA sequence as set with {@link Gene#setSequence(String)}, usually located at [txStart,txEnd) in the corresponding chromosom.
	 * 
	 * @return The gene's DNA sequence as String if set, else null
	 */
	public String getSequence() {
		return sequence;
	}

	/**
	 * <p>Sets the aminoacid sequence this gene is coding for.</p>
	 * 
	 * <p>Should be used carefully because later calculations may depend on the sequence's correctness.</p>
	 * 
	 * @param aaSequence Aminoacid sequence as String
	 */
	public void setAaSequence(String aaSequence) {
		this.aaSequence = aaSequence;
	}

	/**
	 * <p>Returns the aminoacid sequence this gene is coding for, as set with {@link Gene#setAaSequence(String)}.</p>
	 * 
	 * @return Aminoacid sequence as String if set, else null
	 */
	public String getAaSequence() {
		return aaSequence;
	}

	/**
	 * <p>Returns the gene's coding sequence (CDS) in 5'->3' direction, if the DNA sequence is available. 
	 * This is the gene's dna sequence starting with the start codon and usually ending with the stop codon but without introns.</p>
	 * 
	 * @return Coding sequence (CDS), if the DNA sequence is available, null otherwise
	 */
	public String getCds(){	
		if (sequence == null)
			return null;

		int[] starts = getExonStarts().clone();
		int[] ends = getExonEnds().clone();
		int min = 0;
		int max = ends.length;
		while (ends[min] < getCdsStart()){
			min++;
		}
		while (starts[max-1] > getCdsEnd()){
			max--;
		}
		starts[min] = getCdsStart();
		ends[max-1] = getCdsEnd();

		StringBuffer dna = new StringBuffer();

		int start;
		int end;

		for (int i=min; i<max; i++){
			start = starts[i] - txStart;
			end = ends[i] - txStart;
			dna.append(getSequence().substring(start, end));
		}

		return dna.toString().toUpperCase();
	}

	/**
	 * Returns the gene's coding sequence (CDS) length.
	 * 
	 * @return CDS length
	 */
	public int getCdsLength(){	
		int length = 0;

		int[] starts = getExonStarts().clone();
		int[] ends = getExonEnds().clone();
		int min = 0;
		int max = ends.length;
		while (ends[min] < getCdsStart()){
			min++;
		}
		while (starts[max-1] > getCdsEnd()){
			max--;
		}
		starts[min] = getCdsStart();
		ends[max-1] = getCdsEnd();

		for (int i=min; i<max; i++){
			length += ends[i] - starts[i];
		}

		return length;
	}

	/**
	 * Sets the marker "used", which might be useful for expensive operations like loading sequences and keeping them in memory.
	 * 
	 * @param used Set to "true" if used, "false" otherwise
	 */
	public void setUsed(boolean used) {
		this.used = used;
	}

	/**
	 * Returns the value of the "used" marker.
	 * 
	 * @return true or false, depending on how it's set with {@link #setUsed(boolean used)}
	 */
	public boolean isUsed(){
		return used;
	}

	/**
	 * <p>Shifts a cdsStart by a given value, tries to correct txStart and exonStarts near the "new" cdsStart position.
	 * 
	 * <p>If a cdsStart ends up in an intron the nearest exonStart will be extended.
	 * If a cdsStart ends up in an exon (either the original or any other) no exon will be changed.</p>
	 * 
	 * <p>In other words, transcription length is not considered to change the same in length as the cds might do,
	 * except the new cdsStart is inside or in front of the original exon.</p>
	 * 
	 * @param shift Size of the shift
	 */
	public void shiftCdsStart(int shift) {
		// cdsStart adjustment
		cdsStart += shift;
		
		// txStart correction
		if (cdsStart < txStart) txStart = cdsStart;
		
		// exonStarts correction
		for (int i=0; i<exonStarts.length; i++){
			if (cdsStart < exonStarts[i]){
				exonStarts[i] = cdsStart;
				break;
			}
			if (exonStarts[i] <= cdsStart && cdsStart <= exonEnds[i]){
				break;
			}
		}
	}

	/**
	 * <p>Shifts a cdsEnd by a given value, tries to correct txEnd and exonEnds near the "new" cdsStart position.
	 * 
	 * <p>If a cdsEnd ends up in an intron the nearest exonEnd will be extended.
	 * If a cdsEnd ends up in an exon (either the original or any other) no exon will be changed.</p>
	 * 
	 * <p>In other words, transcription length is not considered to change the same in length as the cds might do,
	 * except the new cdsEnd is inside or behind of the original exon.</p>
	 * 
	 * @param shift Size of the shift
	 */
	public void shiftCdsEnd(int shift) {
		// cdsEnd adjustment
		cdsEnd += shift;
		
		// txEnd correction
		if (cdsEnd > txEnd) txEnd = cdsEnd;
		
		// exonEnds correction
		for (int i=exonEnds.length-1; i>=0; i--){
			if (cdsEnd > exonEnds[i]){
				exonEnds[i] = cdsEnd;
				break;
			}
			if (exonStarts[i] < cdsEnd && cdsEnd <= exonEnds[i]){
				break;
			}
		}
	}

	/**
	 * Returns the source for this gene annotation, either Gene.UCSC if the gene is an UCSC knownGene, or Gene.ENSEMBL if it is an Ensemble ensGene.
	 * 
	 * @return String indication the annotation's origin
	 */
	public String getSource() {
		return source;
	}

	/**
	 * Describes the strictness of rejecting imported genes (See {@link #setStrict(boolean)}). 
	 * 
	 * @return true, if genes are strictly rejected
	 */
	public static boolean isStrict() {
		return strict;
	}

	/**
	 * Set the strictness of rejecting imported genes.
	 * 
	 * true : non-coding genes and genes on uncommon chromosomes are rejected.
	 * false: only genes with unrecognized formats (due to regex) are rejected.
	 * 
	 * default is true.
	 * 
	 * @param strict
	 */
	public static void setStrict(boolean strict) {
		Gene.strict = strict;
	}
}
