#! /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.
#
from optparse import OptionParser
from commons.core.parsing.FastaParser import FastaParser
from commons.core.parsing.FastqParser import FastqParser
from commons.core.writer.FastaWriter import FastaWriter
from commons.core.writer.FastqWriter import FastqWriter
from SMART.Java.Python.misc.Progress import Progress
from SMART.Java.Python.misc import Utils


if __name__ == "__main__":
    
    # parse command line
    description = "Trim Sequences v1.0.3: Remove the 5' and/or 3' adaptors of a list of reads. [Category: Data Modification]"

    parser = OptionParser(description = description)
    parser.add_option("-i", "--input",         dest="inputFileName",  action="store",                     type="string", help="input file [compulsory] [format: file in sequence format given by -f]")
    parser.add_option("-f", "--format",        dest="format",         action="store",                     type="string", help="format of file [compulsory] [format: sequence file format]")
    parser.add_option("-o", "--output",        dest="outputFileName", action="store",                     type="string", help="output file [compulsory] [format: output file in sequence format given by -f]")
    parser.add_option("-3", "--threePAdaptor", dest="threePAdaptor",  action="store",      default=None,  type="string", help="3' adaptor [format: string] [default: None]")
    parser.add_option("-5", "--fivePAdaptor",  dest="fivePAdaptor",   action="store",      default=None,  type="string", help="5' adaptor [format: string] [default: None]")
    parser.add_option("-e", "--errors",        dest="errors",         action="store",      default=0,     type="int",    help="number of errors in percent [format: int] [default: 0]")
    parser.add_option("-d", "--indels",        dest="indels",         action="store_true", default=False,                help="also accept indels [format: bool] [default: False]")
    parser.add_option("-n", "--noAdaptor5p",   dest="noAdaptor5p",    action="store",      default=None,  type="string", help="print sequences with no 5' adaptor [format: output file in sequence format given by -f]")
    parser.add_option("-m", "--noAdaptor3p",   dest="noAdaptor3p",    action="store",      default=None,  type="string", help="print sequences with no 3' adaptor [format: output file in sequence format given by -f]")
    parser.add_option("-v", "--verbosity",     dest="verbosity",      action="store",      default=1,     type="int",    help="trace level [format: int]")
    (options, args) = parser.parse_args()

    minSize = 3

    if options.format == "fasta":
        parser = FastaParser(options.inputFileName, options.verbosity)
    elif options.format == "fastq":
        parser = FastqParser(options.inputFileName, options.verbosity)
    else:
        raise Exception("Cannot handle files with '%s' format." % (options.format))

    if options.format == "fasta":
        writer = FastaWriter(options.outputFileName, options.verbosity)
    elif options.format == "fastq":
        writer = FastqWriter(options.outputFileName, options.verbosity)
    else:
        raise Exception("Cannot handle files with '%s' format." % (options.format))


    if options.noAdaptor5p != None:
        if options.format == "fasta":
            writer5pNoAdaptor = FastaWriter(options.noAdaptor5p, options.verbosity)
        elif options.format == "fastq":
            writer5pNoAdaptor = FastqWriter(options.noAdaptor5p, options.verbosity)
        else:
            raise Exception("Cannot handle files with '%s' format." % (options.format))
    nbFound5p = 0
    
    if options.noAdaptor3p != None:
        if options.format == "fasta":
            writer3pNoAdaptor = FastaWriter(options.noAdaptor3p, options.verbosity)
        elif options.format == "fastq":
            writer3pNoAdaptor = FastqWriter(options.noAdaptor3p, options.verbosity)
        else:
            raise Exception("Cannot handle files with '%s' format." % (options.format))
    nbFound3p = 0
            
    progress = Progress(parser.getNbSequences(), "Reading %s" % (options.inputFileName), options.verbosity)
    for sequence in parser.getIterator():
        progress.inc()
        if options.threePAdaptor != None:
            nucleotides = sequence.sequence
            found       = False
            bestScore   = 10000
            bestRegion  = 0
            for i in range(len(nucleotides) - minSize):
                nucleotidesPart = nucleotides[i:]
                adaptorPart     = options.threePAdaptor if len(nucleotidesPart) >= len(options.threePAdaptor) else options.threePAdaptor[:len(nucleotidesPart)]
                nucleotidesPart = nucleotidesPart if len(adaptorPart) == len(nucleotidesPart) else nucleotidesPart[:len(adaptorPart)]
                if options.indels:
                    score = Utils.getLevenshteinDistance(adaptorPart, nucleotidesPart)
                else:
                    score = Utils.getHammingDistance(adaptorPart, nucleotidesPart)
                if score <= int(options.errors / 100.0 * len(adaptorPart)) and score < bestScore:
                    bestScore  = score
                    bestRegion = i
                    found      = True
            if found:
                nbFound3p += 1
                sequence.shrinkToFirstNucleotides(bestRegion)
            elif options.noAdaptor3p:
                writer3pNoAdaptor.addSequence(sequence)
        if options.fivePAdaptor != None:
            nucleotides = sequence.sequence
            found       = False
            bestScore   = 10000
            bestRegion  = 0
            for i in reversed(range(minSize, len(nucleotides))):
                nucleotidesPart = nucleotides[:i]
                adaptorPart     = options.fivePAdaptor if len(nucleotidesPart) >= len(options.fivePAdaptor) else options.fivePAdaptor[-len(nucleotidesPart):]
                nucleotidesPart = nucleotidesPart if len(adaptorPart) == len(nucleotidesPart) else nucleotidesPart[-len(adaptorPart):]
                if options.indels:
                    score = Utils.getLevenshteinDistance(adaptorPart, nucleotidesPart)
                else:
                    score = Utils.getHammingDistance(adaptorPart, nucleotidesPart)
                if score <= int(options.errors / 100.0 * len(adaptorPart)) and score < bestScore:
                    bestScore  = score
                    bestRegion = i
                    found      = True
            if found:
                nbFound5p += 1
                sequence.shrinkToLastNucleotides(len(nucleotides) - bestRegion)
            elif options.noAdaptor5p:
                writer5pNoAdaptor.addSequence(sequence)
        writer.addSequence(sequence)
    progress.done()
    writer.close()

    print "%d sequences" % (parser.getNbSequences())
    if options.fivePAdaptor != None:
        print "%d sequences with 5' adaptors (%.2f%%)" % (nbFound5p, float(nbFound5p) / parser.getNbSequences() * 100)
    if options.threePAdaptor != None:
        print "%d sequences with 3' adaptors (%.2f%%)" % (nbFound3p, float(nbFound3p) / parser.getNbSequences() * 100)