#! /usr/bin/env python
#
# Copyright INRA-URGI 2009-2010
# 
# This software is governed by the CeCILL license under French law and
# abiding by the rules of distribution of free software. You can use,
# modify and/ or redistribute the software under the terms of the CeCILL
# license as circulated by CEA, CNRS and INRIA at the following URL
# "http://www.cecill.info".
# 
# As a counterpart to the access to the source code and rights to copy,
# modify and redistribute granted by the license, users are provided only
# with a limited warranty and the software's author, the holder of the
# economic rights, and the successive licensors have only limited
# liability.
# 
# In this respect, the user's attention is drawn to the risks associated
# with loading, using, modifying and/or developing or reproducing the
# software by the user in light of its specific status of free software,
# that may mean that it is complicated to manipulate, and that also
# therefore means that it is reserved for developers and experienced
# professionals having in-depth computer knowledge. Users are therefore
# encouraged to load and test the software's suitability as regards their
# requirements in conditions enabling the security of their systems and/or
# data to be ensured and, more generally, to use and operate it in the
# same conditions as regards security.
# 
# The fact that you are presently reading this means that you have had
# knowledge of the CeCILL license and that you accept its terms.
#
"""Find random regions in a genome"""

import random, math
from optparse import OptionParser
from commons.core.parsing.FastaParser import *
from commons.core.writer.Gff3Writer import *
from commons.core.writer.MySqlTranscriptWriter import *
from SMART.Java.Python.misc.Progress import *
from SMART.Java.Python.structure.Transcript import Transcript
from SMART.Java.Python.structure.TranscriptContainer import TranscriptContainer

repetitions = 100


class RandomRegionsGenerator(object):

	def __init__(self, verbosity):
		self.verbosity      = verbosity
		self.strands        = False
		self.distribution   = "uniform"
		self.transcripts    = None
		self.sequenceParser = None
		random.seed()


	def setInput(self, fileName):
		self.sequenceParser = FastaParser(fileName, self.verbosity)


	def setGenomeSize(self, size):
		self.genomeSize = size


	def setChromosomeName(self, name):
		self.chromosomeName = name


	def setAnnotation(self, fileName, format):
		parser           = TranscriptContainer(fileName, format, self.verbosity)
		self.transcripts = []
		for transcript in parser.getIterator():
			self.transcripts.append(transcript)
		self.setNumber(len(self.transcripts))
		self.setSize(0)


	def setOutputFile(self, fileName):
		self.outputFileName = fileName


	def setSize(self, size):
		self.minSize = size
		self.maxSize = size


	def setMinSize(self, size):
		self.minSize = size


	def setMaxSize(self, size):
		self.maxSize = size


	def setNumber(self, number):
		self.number = number


	def setStrands(self, strands):
		self.strands = strands


	def setMaxDistribution(self, maxElements):
		if maxElements == None:
			return
		self.maxElements = maxElements
		self.distribution = "gaussian"


	def setDeviationDistribution(self, deviation):
		if deviation == None:
			return
		self.deviation = deviation
		self.distribution = "gaussian"


	def getSizes(self):
		if self.sequenceParser == None:
			self.chromosomes    = [self.chromosomeName]
			self.sizes          = {self.chromosomeName: self.genomeSize}
			self.cumulatedSize  = self.genomeSize
			self.cumulatedSizes = {self.chromosomeName: self.genomeSize}
			return
		self.chromosomes    = self.sequenceParser.getRegions()
		self.sizes          = {}
		self.cumulatedSize  = 0
		self.cumulatedSizes = {}
		for chromosome in self.chromosomes:
			self.sizes[chromosome]          = self.sequenceParser.getSizeOfRegion(chromosome)
			self.cumulatedSize             += self.sizes[chromosome]
			self.cumulatedSizes[chromosome] = self.cumulatedSize


	def findPosition(self, size = None):
		if size == None:
			size = random.randint(self.minSize, self.maxSize)
		integer = random.randint(0, self.cumulatedSize)
		for chromosome in self.chromosomes:
			if self.cumulatedSizes[chromosome] > integer:
				break
		start = random.randint(1, self.sizes[chromosome] - size)
		return (chromosome, start, size)


	def createTranscript(self, chromosome, start, size, strand, cpt):
		transcript = Transcript()
		transcript.setChromosome(chromosome)
		transcript.setEnd(start + size-1)
		transcript.setStart(start)
		transcript.setDirection(strand)
		transcript.setName("rand_%d" % (cpt))
		return transcript


	def moveTranscript(self, chromosome, start, transcript):
		while transcript.getEnd() + start - transcript.getStart() > self.cumulatedSizes[chromosome]:
			chromosome, start, size = self.findPosition(transcript.getEnd() - transcript.getStart())
		newTranscript = Transcript()
		newTranscript.setChromosome(chromosome)
		newTranscript.tags = transcript.tags
		if transcript.getNbExons() > 1:
			for exon in transcript.getNbExons():
				newExon = Interval()
				newExon.setChromosome(chromosome)
				newExon.setEnd(exon.getEnd() + start - transcript.getStart())
				newExon.setStart(exon.getStart() + start - transcript.getStart())
				newTranscript.addExon(newExon)
		newTranscript.setEnd(transcript.getEnd() + start - transcript.getStart())
		newTranscript.setStart(start)
		newTranscript.setDirection(transcript.getDirection())
		return [newTranscript]


	def createUniformCluster(self, chromosome, start, size, strand, cpt):
		transcript = self.createTranscript(chromosome, start, size, strand, cpt)
		return [transcript]


	def findNbTranscripts(self, cpt):
		return min(int(round(math.exp(random.random() * math.log(self.maxElements)))), self.number - cpt + 1)


	def getDev(self):
		deviation = 0.0
		for j in range(repetitions):
			deviation += random.randint(-self.deviation, self.deviation)
		deviation /= repetitions
		deviation  = int(round(deviation))
		return deviation


	def createGaussianCluster(self, chromosome, start, size, strand, cpt):
		transcripts   = []
		nbTranscripts = self.findNbTranscripts(cpt)
		for i in range(nbTranscripts):
			transcript = self.createTranscript(chromosome, start + self.getDev(), size + self.getDev(), strand, cpt + i)
			transcripts.append(transcript)
		return transcripts


	def writeRegions(self):
		writer     = Gff3Writer(self.outputFileName, self.verbosity)
		outputFile = open(self.outputFileName, "w")
		progress   = Progress(self.number, "Writing to %s" % (self.outputFileName), self.verbosity)
		i          = 0
		while i < self.number:
			chromosome, start, size = self.findPosition()
			strand                  = random.choice([-1, 1]) if self.strands else 1
			if self.transcripts != None:
				transcripts = self.moveTranscript(chromosome, start, self.transcripts[i])
			elif self.distribution == "uniform":
				transcripts = self.createUniformCluster(chromosome, start, size, strand, i+1)
			else:
				transcripts = self.createGaussianCluster(chromosome, start, size, strand, i+1)
			for transcript in transcripts:
				writer.addTranscript(transcript)
				i += 1
				progress.inc()
		progress.done()
		outputFile.close()
		writer.write()
		writer.close()


	def run(self):
		self.getSizes()
		self.writeRegions()


if __name__ == "__main__":
	
	# parse command line
	description = "Get Random Regions v1.0.2: Get some random coordinates on a genome. May use uniform or gaussian distribution (in gaussion distribution, # of element per cluster follows a power law). [Category: Other]"

	parser = OptionParser(description = description)
	parser.add_option("-r", "--reference",     dest="reference",      action="store",      default=None,  type="string", help="file that contains the sequences [format: file in FASTA format]")
	parser.add_option("-S", "--referenceSize", dest="referenceSize",  action="store",      default=None,  type="int",    help="size of the chromosome (when no reference is given) [format: int]")
	parser.add_option("-c", "--chromosome",    dest="chromosome",     action="store",      default=None,  type="string", help="name of the chromosome (when no reference is given) [format: string]")
	parser.add_option("-o", "--output",        dest="outputFileName", action="store",                     type="string", help="output file [compulsory] [format: output file in FASTA format]")
	parser.add_option("-i", "--input",         dest="inputFileName",  action="store",      default=None,  type="string", help="optional file containing regions to shuffle [format: file in transcript format given by -f]")
	parser.add_option("-f", "--format",        dest="format",         action="store",      default=None,  type="string", help="format of the previous file [format: transcript file format]")
	parser.add_option("-s", "--size",          dest="size",           action="store",      default=None,  type="int",    help="size of the regions (if no region set is provided) [format: int]")
	parser.add_option("-z", "--minSize",       dest="minSize",        action="store",      default=None,  type="int",    help="minimum size of the regions (if no region set nor a fixed size are provided) [format: int]")
	parser.add_option("-Z", "--maxSize",       dest="maxSize",        action="store",      default=None,  type="int",    help="maximum size of the regions (if no region set nor a fixed size are provided) [format: int]")
	parser.add_option("-n", "--number",        dest="number",         action="store",      default=None,  type="int",    help="number of regions (if no region set is provided) [format: int]")
	parser.add_option("-t", "--strands",       dest="strands",        action="store_true", default=False,                help="use both strands (if no region set is provided) [format: boolean]")
	parser.add_option("-m", "--max",           dest="max",            action="store",      default=None,  type="int",    help="max. # reads in a cluster (for Gaussian dist.) [format: int]")
	parser.add_option("-d", "--deviation",     dest="deviation",      action="store",      default=None,  type="int",    help="deviation around the center of the cluster (for Gaussian dist.) [format: int]")
	parser.add_option("-v", "--verbosity",     dest="verbosity",      action="store",      default=1,     type="int",    help="trace level [format: int]")
	(options, args) = parser.parse_args()

	rrg = RandomRegionsGenerator(options.verbosity)
	if options.reference == None:
		rrg.setGenomeSize(options.referenceSize)
		rrg.setChromosomeName(options.chromosome)
	else:
		rrg.setInput(options.reference)
	rrg.setOutputFile(options.outputFileName)
	if options.inputFileName == None:
		if options.size != None:
			rrg.setSize(options.size)
		else:
			rrg.setMinSize(options.minSize)
			rrg.setMaxSize(options.maxSize)
		rrg.setNumber(options.number)
		rrg.setStrands(options.strands)
	else:
		rrg.setAnnotation(options.inputFileName, options.format)
	rrg.setMaxDistribution(options.max)
	rrg.setDeviationDistribution(options.deviation)
	rrg.run()