#!/usr/bin/env python
# reverse transcribe a peptide alignment into a codon aligned alignment


import optparse
import os
from rasmus import util
from compbio import fasta, alignlib


#=============================================================================
o = optparse.OptionParser()
o.add_option("-d", "--dna", dest="dna", metavar="<dna fasta>",
             action="append")
o.add_option("-o", "--oldext", dest="oldext", metavar="<old extension>",
             default=".align")
o.add_option("-n", "--newext", dest="newext", metavar="<new extension>",
             default=".nt.align")
o.add_option("-i", "--useindex", dest="useindex", action="store_true",
             default=False)
o.add_option("-t", "--transmap", dest="transmap",
             help="mapping of transcript names to proteins")
o.add_option("-c", "--check", dest="check", action="store_true", default=False,
             help="check that the dna translates to the peptide")
o.add_option("--trim", action="store_true")


conf, args = o.parse_args()

#=============================================================================


class SeqMap (object):

    def __init__(self, seqs, transmap):
        self.seqs = seqs
        self.transmap = transmap

    def __getitem__(self, key):
        return self.seqs[self.transmap[key]]


seqs = fasta.FastaDict()

# read dna seqs
util.tic("read DNA sequences")
for f in conf.dna:
    util.tic("read fasta '%s'" % f)
    seqs.read(f, useIndex=conf.useindex)
    util.toc()
util.toc()

if conf.transmap:
    transmap = {}
    for row in util.DelimReader(conf.transmap):
        transmap[row[0]] = row[1]
    seqs = SeqMap(seqs, transmap)


# process alignments
util.tic("process alignments")
for aln_file in args:
    if not os.path.exists(aln_file):
        util.log("skipping '%s', does not exist" % aln_file)
        continue

    newfile = util.replace_ext(aln_file, conf.oldext, conf.newext)
    util.log(aln_file, "===>", newfile)
    aln = fasta.read_fasta(aln_file)

    try:
        alndna = alignlib.revtranslate_align(aln, seqs,
                                             check=conf.check,
                                             trim=conf.trim)
    except alignlib.TranslateError, e:
        print e
        print "%s  ! ===> %s" % (e.a, e.codon)
        print e.dna
        print e.aa
    else:
        alndna.write(newfile)

util.toc()
