#! /usr/bin/env python

 # Copyright 2014 Martin C. Frith

 # Read MAF-format alignments, and write those that have a segment with

 # score >= threshold, with gentle masking of lowercase letters.  There

 # must be a lastal header with score parameters.

 # Gentle masking is described in: MC Frith, PLoS One 2011;6(12):e28819

 # "Gentle masking of low-complexity sequences improves homology search"

 # Limitations: doesn't (yet) handle sequence quality data,

 # frameshifts, or generalized affine gaps.

 import itertools, optparse, os, signal, sys

 def getScoreMatrix(rowHeads, colHeads, matrix, deleteCost, insertCost):

     defaultScore = min(map(min, matrix))

     scoreMatrix = [[defaultScore for i in range(128)] for j in range(128)]

     for i, x in enumerate(rowHeads):

         for j, y in enumerate(colHeads):

             xu = ord(x.upper())

             xl = ord(x.lower())

             yu = ord(y.upper())

             yl = ord(y.lower())

             score = matrix[i][j]

             maskScore = min(score, 0)

             scoreMatrix[xu][yu] = score

             scoreMatrix[xu][yl] = maskScore

             scoreMatrix[xl][yu] = maskScore

             scoreMatrix[xl][yl] = maskScore

     for i in range(128):

         scoreMatrix[i][ord("-")] = -deleteCost

         scoreMatrix[ord("-")][i] = -insertCost

     return scoreMatrix

 def isGoodAlignment(seqs, scoreMatrix, delOpenCost, insOpenCost, minScore):

     """Does the alignment have a segment with score >= minScore?"""

     r, q = seqs

     score = 0

     xOld = " "

     yOld = " "

     for x, y in itertools.izip(r, q):

         score += scoreMatrix[ord(x)][ord(y)]

         if score >= minScore: return True

         if x == "-" and xOld != "-": score -= insOpenCost

         if y == "-" and yOld != "-": score -= delOpenCost

         if score < 0: score = 0

         xOld = x

         yOld = y

     return False

 def printIfGood(maf, seqs, scoreMatrix, delOpenCost, insOpenCost, minScore):

     if isGoodAlignment(seqs, scoreMatrix, delOpenCost, insOpenCost, minScore):

         for line in maf:

             print line,

         print

 def doOneFile(lines):

     scoreMatrix = []

     maf = []

     seqs = []

     for line in lines:

         if line[0] == "#":

             print line,

             w = line.split()

             for i in w:

                 if i.startswith("a="): aDel = int(i[2:])

                 if i.startswith("b="): bDel = int(i[2:])

                 if i.startswith("A="): aIns = int(i[2:])

                 if i.startswith("B="): bIns = int(i[2:])

                 if i.startswith("e="): minScore = int(i[2:])

             if len(w) > 1 and max(map(len, w)) == 1:

                 colHeads = w[1:]

                 rowHeads = []

                 matrix = []

             elif len(w) > 2 and len(w[1]) == 1:

                 rowHeads.append(w[1])

                 matrix.append(map(int, w[2:]))

         elif line.isspace():

             if seqs: printIfGood(maf, seqs, scoreMatrix, aDel, aIns, minScore)

             maf = []

             seqs = []

         else:

             if not scoreMatrix:

                 scoreMatrix = getScoreMatrix(rowHeads, colHeads, matrix,

                                              bDel, bIns)

             maf.append(line)

             if line[0] == "s": seqs.append(line.split()[6])

     if seqs: printIfGood(maf, seqs, scoreMatrix, aDel, aIns, minScore)

 def lastPostmask(args):

     if args:

         for i in args:

             if i == "-":

                 doOneFile(sys.stdin)

             else:

                 with open(i) as f:

                     doOneFile(f)

     else:

         doOneFile(sys.stdin)

 if __name__ == "__main__":

     signal.signal(signal.SIGPIPE, signal.SIG_DFL)  # avoid silly error message

     usage = "%prog in.maf > out.maf"

     description = "Get alignments that have a segment with score >= threshold, with gentle masking of lowercase letters."

     op = optparse.OptionParser(usage=usage, description=description)

     (opts, args) = op.parse_args()

     try: lastPostmask(args)

     except KeyboardInterrupt: pass  # avoid silly error message

     except Exception, e:

         prog = os.path.basename(sys.argv[0])

         sys.exit(prog + ": error: " + str(e))