#!/usr/bin/env bds

#-------------------------------------------------------------------------------
#
# Download genome data
#
#-------------------------------------------------------------------------------

ucscMysqlServer    := "genome-mysql.cse.ucsc.edu"
cse                := "hgdownload.cse.ucsc.edu"
soe                := "hgdownload.soe.ucsc.edu"
urlUcscCse         := "http://$cse/goldenPath"
urlUcscSoe         := "http://$soe/goldenPath"
ftpUcscSoe         := "ftp://$soe/goldenPath"

void downloadUcscGenome(string genome) {
	if( genome.endsWith('kg') ) {
		dg := new DownloadUcscKg(genome)
		dg.download()
	} else {
		dg := new DownloadUcsc(genome)
		dg.download()
	}
}


# Class to download files and query MySql database from UCSC
class DownloadUcsc extends DownloadGenome {
	string fastaTarUrl
	string cytoBandUrl
	string cdsUrlAlt
	string genesUrlAlt

	void DownloadUcsc(string genome) {
		super.DownloadGenome(genome)
		cdsUrl = "$urlUcscSoe/$genome/bigZips/refMrna.fa.gz"
		cdsUrlAlt = "$urlUcscSoe/$genome/bigZips/mrna.fa.gz"
		fastaUrl = "$urlUcscCse/$genome/bigZips/$genome\.fa.gz"
		fastaTarUrl = "$urlUcscCse/$genome/bigZips/chromFa.tar.gz"
		genesUrl = "$urlUcscSoe/$genome/bigZips/genes/$genome\.refGene.gtf.gz"
		genesUrlAlt = "$urlUcscSoe/$genome/bigZips/genes/$genome\.ncbiRefSeq.gtf.gz"
		cytoBandUrl = "$urlUcscCse/$genome/database/cytoBand.txt.gz"
		protUrl = ''  # It looks like there is no protein file with RefSeq transcripts at this moment
	}

	string downloadCdsFasta() {
	    if( cdsUrl.exists() ) return downloadFile(cdsUrl, cdsLocal, "CDS FASTA", false)
	    if( cdsUrlAlt.exists() ) return downloadFile(cdsUrlAlt, cdsLocal, "CDS FASTA", false)
	    error "Could not find CDS / mRNA FASTA file"
	}

	string downloadCytoBand() {
		cytoBandFile := "$downloadDir/cytoBand.txt.gz"
		if( !cytoBandFile.exists() && cytoBandUrl.exists() ) {
			wget(cytoBandUrl, cytoBandFile)
			return cytoBandFile
		}
		return '';
	}

	# Download FASTA geome from UCSC
	string downloadFasta() {
		if( fastaUrl.exists() ) {
			wgetTask(fastaUrl, fastaLocal)
		} else if( fastaTarUrl.exists() ) {
			localTar := "$downloadDir/chromFa.tar.gz"
			wgetTask(fastaTarUrl, localTar)
			fastaLocal = fastaFromTar(localTar) # Create fasta from tar
		} else {
			println "Could not download FASTA reference genome file\n\tFASTA: '$fastaUrl'\n\tTAR: '$fastaTarUrl'"
		}

		return fastaLocal
	}

	# Download genes file. Either download GTF or query MySql server
	string downloadGenes() {
		genes := downloadGtf()
		if( genes.isEmpty() ) {
			# GTF not found, let's query MySql database
			genes = query()
		}
		return genes
	}

	# Download geome from UCSC, ony if it has a GTF file
	string downloadGtf() {
		log("Downloading '$genome' (UCSC): GTF file")
		if( genesUrl.exists() ) return wgetTask(genesUrl, genesLocal) ? genesLocal : ''
		if( genesUrlAlt.exists() ) return wgetTask(genesUrlAlt, genesLocal) ? genesLocal : ''
		error "Could not download genes from '$genesUrl' or '$genesUrlAlt'"
	}

	string downloadProteinFasta() {
		# We don't get protein fasta files from UCSC: Nothing to do
		return ''
	}

	# Download other files
	string[] downloadOther() {
		string[] files
		cyto := downloadCytoBand()
		if( !cyto.isEmpty() ) files += cyto
		return files
	}

	# Create a single fasta file from a set of fasta in a tar
	string fastaFromTar(string tar) {
		log("Creating '$genome' FASTA file from tarball")
		fasta := "$downloadDir/$genome\.fa.gz"
		dirChr := "$downloadDir/chr"
		dirChr.mkdir()

		task( fasta <- tar ) {
			sys echo "Creating fasta file '$fasta' from tar file '$tar'"
			sys cd '$dirChr'
			sys tar -xvzf '$tar'
			sys rm -f '$fasta'
			sys cat `find . -iname "*.fa"` | gzip -c > 'fasta'
		}

		return fasta
	}

	# Download mRNA file and create a 'cds.fa.gz'
	string processCds(string cdsLocal) {
		cds := "$downloadDir/cds.fa.gz"
		task(cds <- cdsLocal) {
			sys echo 'Creating CDS FASTA from "$cdsLocal"'
			sys gunzip -c '$cdsLocal' | tr " " "." | gzip -c > '$cds'
		}
		return cds
	}

	string query() {
		# Query gene database
		if( ucscGenomesRefGene.has(genome) ) return queryUcscRefGene()
		return queryUcscXenoRefGene()
	}

	# Query MySql's database (RefSeq)
	string queryUcscRefGene() {
		refSeq := "$downloadDir/genes.refseq.gz"
		if( !refSeq.exists() ) {
			log("Querying UCSC database for RefSeq genes: $genome")
			query := "
				use $genome;
				select rg.bin as 'binNum'
						, CONCAT(rg.name,'.',gi.version) as 'name'
						, rg.chrom
						, rg.strand
						, rg.txStart
						, rg.txEnd
						, rg.cdsStart
						, rg.cdsEnd
						, rg.exonCount
						, rg.exonStarts
						, rg.exonEnds
						, rg.score
						, rg.name2
						, rg.cdsStartStat
						, rg.cdsEndStat
						, rg.exonFrames
				from refGene rg
					inner join hgFixed.gbCdnaInfo gi
					on rg.name=gi.acc
				;
				"
			sys echo "$query" | mysql --user=genome --host=$ucscMysqlServer -A $genome | gzip -c > '$refSeq'
		}
		return refSeq
	}

	# Query MySql's database (dump xenoRefGene table)
	string queryUcscXenoRefGene() {
		refSeq := "$downloadDir/genes.refseq\.gz"
		if( !refSeq.exists() ) {
			log("Querying database for xenoRefGene genes: $genome")
			query := "
				use $genome;
				select rg.bin as 'binNum'
						, rg.name as 'name'
						, rg.chrom
						, rg.strand
						, rg.txStart
						, rg.txEnd
						, rg.cdsStart
						, rg.cdsEnd
						, rg.exonCount
						, rg.exonStarts
						, rg.exonEnds
						, rg.score
						, rg.name2
						, rg.cdsStartStat
						, rg.cdsEndStat
						, rg.exonFrames
				from xenoRefGene rg
				;
				"
			sys echo "$query" | mysql --user=genome --host=$ucscMysqlServer -A $genome | gzip -c > '$refSeq' || true
		}
		return refSeq
	}
}


# Download a 'knaownGenes' version of the genome
class DownloadUcscKg extends DownloadUcsc {
	string refdb

	void DownloadUcscKg(string genome) {
		super.DownloadUcsc(genome)
		this.refdb = genome.substr(0, genome.length() - 2)	# Remove the 'kg' trailing

		genesUrl = "$urlUcscSoe/$this.genome/bigZips/genes/$this.refdb\.knownGene.gtf.gz"
		cdsUrl = ''   # It looks like there is no CDS file with knownGene transcripts at this moment
		protUrl = ''  # It looks like there is no protein file with knownGene transcripts at this moment
		cytoBandUrl = "$urlUcscCse/$refdb/database/cytoBand.txt.gz"

		fastaDataEq = "$dirData/genomes/$refdb\.fa.gz"
	}

	string query() {
		return queryUcscKg()
	}

	# Query MySql's database (KnonwGene)
	string queryUcscKg() {
		refSeq := "$downloadDir/genes.kg.gz"
		if( !refSeq.exists() ) {
			log("Querying database for KnownGenes genes: $genome")
			query := "
				use $genome;
				select kg.name as '#name'
						, kg.chrom
						, kg.strand
						, kg.txStart
						, kg.txEnd
						, kg.cdsStart
						, kg.cdsEnd
						, kg.exonCount
						, kg.exonStarts
						, kg.exonEnds
						, kg.proteinID
						, kg.alignID
				from knownGene kg
				;
				"
			sys echo "$query" | mysql --user=genome --host=$ucscMysqlServer -A $genome | gzip -c > '$refSeq'
		}
		return refSeq
	}
}
