#!/usr/bin/env python
"""Code for performing alignments by Needleman-Wunsch and Smith-Waterman.
"""

__author__ = "Rob Knight"
__copyright__ = "Copyright 2007-2009, The Cogent Project"
__credits__ = ["Rob Knight", "Jeremy Widmann"]
__license__ = "GPL"
__version__ = "1.4.1"
__maintainer__ = "Rob Knight"
__email__ = "rob@spot.colorado.edu"
__status__ = "Development"

class ScoreCell(object):
    """Cell in a ScoreMatrix object. Contains score and pointer."""
    __slots__ = ['Score', 'Pointer']
    def __init__(self, Score=0, Pointer=None):
        """Returns new ScoreCell object.

        Score should be a number.
        Pointer should be a direction (typically None, "up", "diag", or "left")
        """
        self.Score = Score
        self.Pointer = Pointer

    def update(self, diag, up, left):
        """Updates the ScoreCell object with scores for its three neighbors.

        Finds max of (diag, up, and left), and sets the score and pointer
        accordingly. In case of tie, picks up (i.e. tries to minimize
        gaps since the first sequence is the shortest), then diag, then left.
        """
        best = max(diag, up, left)
        if best == up:
            self.Score = up
            self.Pointer = "up"
        elif best == diag:
            self.Score = diag
            self.Pointer = "diag"
        else:
            self.Score = left
            self.Pointer = "left"

def MatchScorer(match, mismatch):
    """Factory function that returns a score function set to match and mismatch.

    match and mismatch should both be numbers. Typically, match should be 
    positive and mismatch should be negative.

    Resulting function has signature f(x,y) -> number.
    """
    def scorer(x, y):
        if x == y:
            return match
        else:
            return mismatch
    return scorer

equality_scorer = MatchScorer(1, -1)
default_gap = -1
default_gap_symbol = '-'

class ScoreMatrix(list):
    """Matrix that contains (score, pointer) pairs for sequence alignment."""
    
    def __init__(self, First, Second, Score=equality_scorer, \
        GapScore=default_gap, GapSymbol=default_gap_symbol):
        """Returns new ScoreMatrix object, initialized but not filled.

        Calls internal methods to initialize first row, column and cell.

        First and Second are the two sequences that will be aligned. Columns
        correspond to positions in First; rows correspond to positions in 
        Second.
        
        Score is a scoring function that takes two corresponding items in
        First and Second, and returns a number that can be compared and added.

        Gap is the score for inserting a gap.

        Note that the matrix is one row and one column bigger than the lengths
        of the sequences, since the first row and first column contain 
        initialization data.
        """
        self.First = First          #first sequence
        self.Second = Second        #second sequence
        self.Cols = len(First) + 1  #need to add 1 for the start col
        self.Rows = len(Second) + 1 #need to add 1 for the start row
        self.Score = Score          #scoring function: f(x,y) = num
        self.GapScore = GapScore    #gap penalty
        self.GapSymbol = GapSymbol  #symbol for gaps
        self.Filled = False         #whether the matrix has been filled
        self.FirstAlign = []        #first sequence, aligned, as list
        self.SecondAlign = []       #second sequence, aligned, as list
        for r in range(self.Rows):
            self.append([ScoreCell() for c in range(self.Cols)])
        self._init_first_row()
        self._init_first_col()
        self._init_first_cell()
        #print "INITIALIZED WITH:"
        #print self

    def __str__(self):
        """Prints score matrix, including labels for row and column."""
        first = self.First
        second = self.Second
        if not first or not second:
            return "Empty Score Matrix"
        #add first sequence as first row: skip 2 positions for row label and
        #for the first column, which is initialized to default values
        rows = ['\t\t' + '\t'.join(self.First)]
        for index, row in enumerate(self):
            #if it's not the first row, add the appropriate char from the seq
            if index:
                curr = second[index-1]
            else:
                curr = ''
            #add the row containing the char and the data
            rows.append('%s\t'%curr + '\t'.join([str(i.Score) for i in row]))
        return '\n'.join(rows)

    def _init_first_row(self):
        """Hook for matrices that need to initialize the first row."""
        pass

    def _init_first_col(self):
        """Hook for matrices that need to initialize the first column."""
        pass

    def _init_first_cell(self):
        """Hook for matrices that need to initialize the first cell."""
        pass

    def alignment(self):
        """Returns alignment of first and second sequences.

        Calls fill() and traceback() if necessary.
        
        Converts the aligned versions to the same class as the originals.
        """
        seq1, seq2 = self.First, self.Second
        aln1, aln2 = self.FirstAlign, self.SecondAlign
        if not aln1 or not aln2:
            self.fill()
            self.traceback()
            aln1, aln2 = self.FirstAlign, self.SecondAlign
        #Remember to return sequences that are the correct class. If the class
        #subclasses str, you probably want ''.join rather than str to feed into
        #the constructor, since str() on a list prints the brackets and commas.
        if isinstance(seq1, str):
            first_result = seq1.__class__(''.join(aln1))
        else:
            first_result = seq1.__class__(aln1)
            
        if isinstance(seq2, str):
            second_result = seq2.__class__(''.join(aln2))
        else:
            second_result = seq2.__class__(aln2)

        return first_result, second_result
    

class NeedlemanWunschMatrix(ScoreMatrix):
    """Score matrix for global alignment using Needleman-Wunsch."""

    def _init_first_row(self):
        """First row initialized with index * gap penalty."""
        gap = self.GapScore
        self[0] = [ScoreCell(gap * i, 'left') for i in range(self.Cols)]
        #print "AFTER FIRST ROW:"
        #print self

    def _init_first_col(self):
        """First column initialized with index * gap penalty."""
        gap = self.GapScore
        for index, row in enumerate(self):
            row[0] = ScoreCell(gap * index, 'up')
        #print "AFTER FIRST COL:"
        #print self

    def _init_first_cell(self):
        """First cell initialized to 0 with no pointer."""
        self[0][0] = ScoreCell(0)
        #print "AFTER FIRST CELL:"
        #print self

    def fill(self):
        """Fills each cell with its best score and the direction of the next.

        For moving up or left, calculates the gap penalty plus the score of the
        cell. For moving diagonally, calculates the match/mismatch score for
        the corresponding positions in the sequence plus the score of the cell.

        Performs all calculations in place.
        """
        score = self.Score
        gap = self.GapScore
        seq_1, seq_2 = self.First, self.Second
        curr_row = self[0]
        for row_i in range(1, self.Rows):
            prev_row = curr_row
            curr_row = self[row_i]
            curr_cell = curr_row[0]
            for col_i in range(1, self.Cols):
                prev_cell = curr_cell
                curr_cell = curr_row[col_i]
                #remember to subtract 1 to find corresponding pos in seq
                diag_score = score(seq_1[col_i-1], seq_2[row_i-1]) + \
                    prev_row[col_i-1].Score
                up_score = gap + prev_row[col_i].Score
                left_score = gap + prev_cell.Score
                #print 'row: %s col: %s '%(row_i,col_i), \
                #    diag_score,up_score,left_score
                curr_cell.update(diag_score, up_score, left_score)
        self.Filled = True
        self.MaxScore = (self[-1][-1].Score, len(self)-1, len(self[0])-1)
        #print "FINISHED FILL"
        #print self

    def traceback(self):
        """Returns the optimal alignment as a (first, second) tuple w/ gaps.

        Follows the pointers assigned in fill(), inserting gaps whenever the
        movement is not diagonal.
        """
        if not self.Filled:
            self.fill()
        align_1 = []
        align_2 = []
        seq_1 = self.First
        seq_2 = self.Second
        curr_row = self.Rows - 1    #subtract 1 to use length as index
        curr_col = self.Cols - 1    #subtract 1 to use length as index
        p = self[curr_row][curr_col].Pointer
        while 1:
            #print "ROW: %s, COL: %s" % (curr_row, curr_col),p
            if p == 'diag':
                align_1.append(seq_1[curr_col-1])
                align_2.append(seq_2[curr_row-1])
                curr_row -= 1
                curr_col -= 1
            elif p == 'left':
                align_1.append(seq_1[curr_col-1])
                align_2.append('-')
                curr_col -= 1
            elif p == 'up':
                align_1.append('-')
                align_2.append(seq_2[curr_row-1])
                curr_row -= 1
            else:
                break
            p = self[curr_row][curr_col].Pointer
        align_1.reverse()
        align_2.reverse()
        self.FirstAlign, self.SecondAlign = align_1, align_2

class SmithWatermanMatrix(ScoreMatrix):

    def fill(self):
        max_row, max_col, max_score = 0, 0, 0
        score = self.Score
        gap = self.GapScore
        seq_1, seq_2 = self.First, self.Second
        curr_row = self[0]
        for row_i in range(1, self.Rows):
            prev_row = curr_row
            curr_row = self[row_i]
            curr_cell = curr_row[0]
            for col_i in range(1, self.Cols):
                prev_cell = curr_cell
                curr_cell = curr_row[col_i]
                #remember to subtract 1 to find corresponding pos in seq
                diag_score = score(seq_1[col_i-1], seq_2[row_i-1]) + \
                    prev_row[col_i-1].Score
                up_score = gap + prev_row[col_i].Score
                left_score = gap + prev_cell.Score
                if max(up_score, left_score, diag_score) <= 0:
                    continue #leave uninitialized if scores all below threshold
                curr_cell.update(diag_score, up_score, left_score)
                if curr_cell.Score > max_score:
                    max_score = curr_cell.Score
                    max_row = row_i
                    max_col = col_i
        self.MaxScore = (max_score, max_row, max_col)
                    
        #print "FINISHED FILL"
        #print self


    def traceback(self):
        align_1 = []
        align_2 = []
        seq_1 = self.First
        seq_2 = self.Second
        max_score, curr_row, curr_col = self.MaxScore
        p = self[curr_row][curr_col].Pointer
        while 1:
            #print "ROW: %s, COL: %s" % (curr_row, curr_col)
            if p == 'diag':
                align_1.append(seq_1[curr_col-1])
                align_2.append(seq_2[curr_row-1])
                curr_row -= 1
                curr_col -= 1
            elif p == 'left':
                align_1.append(seq_1[curr_col-1])
                align_2.append('-')
                curr_col -= 1
            elif p == 'up':
                align_1.append('-')
                align_2.append(seq_2[curr_row-1])
                curr_row -= 1
            else:
                break
            p = self[curr_row][curr_col].Pointer
        align_1.reverse()
        align_2.reverse()
        self.FirstAlign, self.SecondAlign = align_1, align_2

def nw_align(seq1, seq2, scorer=equality_scorer, gap=default_gap, return_score=False):
    """Returns globally optimal alignment of seq1 and seq2."""
    N = NeedlemanWunschMatrix(seq1, seq2, scorer, gap)
    if return_score:
        return N.alignment(), N.MaxScore[0]
    else:
        return N.alignment()

def sw_align(seq1, seq2, scorer=equality_scorer, gap=default_gap, return_score=False):
    """Returns locally optimal alignment of seq1 and seq2."""
    S = SmithWatermanMatrix(seq1, seq2, scorer, gap)
    if return_score:
        return S.alignment(), S.MaxScore[0]
    else:
        return S.alignment()

def demo(seq1, seq2):
    result = []
    result.append("Global alignment:")
    result.extend(nw_align(seq1, seq2))
    result.append("Local alignment:")
    result.extend(sw_align(seq1, seq2))
    return "\n".join(result)

#initialization

if __name__ == '__main__':
    from sys import argv
    print demo(argv[1], argv[2])