#! /usr/bin/env python
# Copyright 2010, 2011, 2013, 2014 Martin C. Frith
# Read MAF-format alignments: write them in other formats.
# Seems to work with Python 2.x, x>=4

# By "MAF" we mean "multiple alignment format" described in the UCSC
# Genome FAQ, not e.g. "MIRA assembly format".

from itertools import *
import sys, os, fileinput, math, operator, optparse, signal, string

def maxlen(s):
    return max(map(len, s))

def joined(things, delimiter):
    return delimiter.join(map(str, things))

identityTranslation = string.maketrans("", "")
def deleted(myString, deleteChars):
    return myString.translate(identityTranslation, deleteChars)

def quantify(iterable, pred=bool):
    """Count how many times the predicate is true."""
    return sum(map(pred, iterable))

class Maf:
    def __init__(self, aLine, sLines, qLines, pLines):
        self.namesAndValues = dict(i.split("=") for i in aLine[1:])

        if not sLines: raise Exception("empty alignment")
        cols = zip(*sLines)
        self.seqNames = cols[1]
        self.alnStarts = map(int, cols[2])
        self.alnSizes = map(int, cols[3])
        self.strands = cols[4]
        self.seqSizes = map(int, cols[5])
        self.alnStrings = cols[6]

        self.qLines = qLines
        self.pLines = pLines

def dieUnlessPairwise(maf):
    if len(maf.alnStrings) != 2:
        raise Exception("pairwise alignments only, please")

def insertionCounts(alnString):
    gaps = alnString.count("-")
    forwardFrameshifts = alnString.count("\\")
    reverseFrameshifts = alnString.count("/")
    letters = len(alnString) - gaps - forwardFrameshifts - reverseFrameshifts
    return letters, forwardFrameshifts, reverseFrameshifts

def coordinatesPerLetter(alnString, alnSize):
    letters, forwardShifts, reverseShifts = insertionCounts(alnString)
    if forwardShifts or reverseShifts or letters < alnSize: return 3
    else: return 1

def mafLetterSizes(maf):
    return map(coordinatesPerLetter, maf.alnStrings, maf.alnSizes)

def alnLetterSizes(sLines):
    return [coordinatesPerLetter(i[6], int(i[3])) for i in sLines]

def isTranslated(sLines):
    for i in sLines:
        alnString = i[6]
        if "\\" in alnString or "/" in alnString: return True
        if len(alnString) - alnString.count("-") < int(i[3]): return True
    return False

def insertionSize(alnString, letterSize):
    """Get the length of sequence included in the alnString."""
    letters, forwardShifts, reverseShifts = insertionCounts(alnString)
    return letters * letterSize + forwardShifts - reverseShifts

def symbolSize(symbol, letterSize):
    if symbol == "-": return 0
    if symbol == "\\": return 1
    if symbol == "/": return -1
    return letterSize

def isMatch(alignmentColumn):
    # No special treatment of ambiguous bases/residues: same as NCBI BLAST.
    first = alignmentColumn[0].upper()
    for i in alignmentColumn[1:]:
        if i.upper() != first: return False
    return True

def isGapless(alignmentColumn):
    return "-" not in alignmentColumn

def matchAndInsertSizes(alignmentColumns, letterSizes):
    """Get sizes of gapless blocks, and of the inserts between them."""
    for k, v in groupby(alignmentColumns, isGapless):
        if k:
            sizeOfGaplessBlock = len(list(v))
            yield str(sizeOfGaplessBlock)
        else:
            blockRows = izip(*v)
            blockRows = imap(''.join, blockRows)
            insertSizes = imap(insertionSize, blockRows, letterSizes)
            yield joined(insertSizes, ":")

##### Routines for converting to AXT format: #####

axtCounter = count()

def writeAxt(maf):
    if maf.strands[0] != "+":
        raise Exception("for AXT, the 1st strand in each alignment must be +")

    # Convert to AXT's 1-based coordinates:
    alnStarts = imap(operator.add, maf.alnStarts, repeat(1))
    alnEnds = imap(operator.add, maf.alnStarts, maf.alnSizes)

    rows = zip(maf.seqNames, alnStarts, alnEnds, maf.strands)
    head, tail = rows[0], rows[1:]

    outWords = []
    outWords.append(axtCounter.next())
    outWords.extend(head[:3])
    outWords.extend(chain(*tail))
    outWords.append(maf.namesAndValues["score"])

    print joined(outWords, " ")
    for i in maf.alnStrings: print i
    print  # print a blank line at the end

##### Routines for converting to tabular format: #####

def writeTab(maf):
    aLine, sLines, qLines, pLines = maf

    score = "0"
    expect = None
    mismap = None
    for i in aLine:
        if   i.startswith("score="): score = i[6:]
        elif i.startswith("expect="): expect = i[7:]
        elif i.startswith("mismap="): mismap = i[7:]

    outWords = []
    outWords.append(score)

    for i in sLines: outWords.extend(i[1:6])

    alnStrings = [i[6] for i in sLines]
    alignmentColumns = zip(*alnStrings)
    letterSizes = alnLetterSizes(sLines)
    gapStrings = matchAndInsertSizes(alignmentColumns, letterSizes)
    gapWord = ",".join(gapStrings)
    outWords.append(gapWord)

    if expect: outWords.append(expect)
    if mismap: outWords.append(mismap)

    print "\t".join(outWords)

##### Routines for converting to PSL format: #####

def pslBlocks(alignmentColumns, alnStarts, letterSizes):
    """Get sizes and start coordinates of gapless blocks in an alignment."""
    start1, start2 = alnStarts
    letterSize1, letterSize2 = letterSizes
    size = 0
    for x, y in alignmentColumns:
        if x == "-":
            if size:
                yield size, start1, start2
                start1 += size * letterSize1
                start2 += size * letterSize2
                size = 0
            start2 += symbolSize(y, letterSize2)
        elif y == "-":
            if size:
                yield size, start1, start2
                start1 += size * letterSize1
                start2 += size * letterSize2
                size = 0
            start1 += symbolSize(x, letterSize1)
        else:
            size += 1
    if size: yield size, start1, start2

def pslCommaString(things):
    # UCSC software seems to prefer a trailing comma
    return joined(things, ",") + ","

def gapRunCount(letters):
    """Get the number of runs of gap characters."""
    uniqLetters = map(operator.itemgetter(0), groupby(letters))
    return uniqLetters.count("-")

def pslEndpoints(alnStart, alnSize, strand, seqSize):
    alnEnd = alnStart + alnSize
    if strand == "+": return alnStart, alnEnd
    else: return seqSize - alnEnd, seqSize - alnStart

def caseSensitivePairwiseMatchCounts(columns, isProtein):
    # repMatches is always zero
    # for proteins, nCount is always zero, because that's what BLATv34 does
    standardBases = "ACGTU"
    matches = mismatches = repMatches = nCount = 0
    for i in columns:
        if "-" in i: continue
        x, y = i
        if x in standardBases and y in standardBases or isProtein:
            if x == y: matches += 1
            else: mismatches += 1
        else: nCount += 1
    return matches, mismatches, repMatches, nCount

def writePsl(maf, isProtein):
    dieUnlessPairwise(maf)

    alnStrings = map(str.upper, maf.alnStrings)
    alignmentColumns = zip(*alnStrings)
    letterSizes = mafLetterSizes(maf)

    matchCounts = caseSensitivePairwiseMatchCounts(alignmentColumns, isProtein)
    matches, mismatches, repMatches, nCount = matchCounts
    numGaplessColumns = sum(matchCounts)

    qNumInsert = gapRunCount(maf.alnStrings[0])
    qBaseInsert = maf.alnSizes[1] - numGaplessColumns * letterSizes[1]

    tNumInsert = gapRunCount(maf.alnStrings[1])
    tBaseInsert = maf.alnSizes[0] - numGaplessColumns * letterSizes[0]

    strand = maf.strands[1]
    if max(letterSizes) > 1:
        strand += maf.strands[0]
    elif maf.strands[0] != "+":
        raise Exception("for non-translated PSL, the 1st strand in each alignment must be +")

    tName, qName = maf.seqNames
    tSeqSize, qSeqSize = maf.seqSizes

    rows = zip(maf.alnStarts, maf.alnSizes, maf.strands, maf.seqSizes)
    tStart, tEnd = pslEndpoints(*rows[0])
    qStart, qEnd = pslEndpoints(*rows[1])

    blocks = list(pslBlocks(alignmentColumns, maf.alnStarts, letterSizes))
    blockCount = len(blocks)
    blockSizes, tStarts, qStarts = map(pslCommaString, zip(*blocks))                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                        