# nw.py
#   A simple implementation of the Needleman-Wunsch algorithm,
#   intended as the basis for exercises in a Bioinformatics course.
# Author: Samuel A. Rebelsky

import string

# +--------+----------------------------------------------------------
# | Values |
# +--------+

# A dictionary that gives the 'value' of different actions.  We use
# this dictonary in the computation of the Needleman-Wunsch table.
# The dictionary should map the values 'Del', 'Ins', 'Mat', and 'Mis'
# (for 'Delete', 'Insert', 'Match', or 'Mismatch') to the values
# associated with those different actions.
actionValuesDefault = {}
actionValuesDefault['Del'] = -1
actionValuesDefault['Ins'] = -1
actionValuesDefault['Mat'] = 1
actionValuesDefault['Mis'] = 0


# +------------+------------------------------------------------------
# | Procedures |
# +------------+

def read_fasta(filename):
    """Read fasta data from the given file.  Returns a two-element list,
       the first of which is the fasta information (the first line), the
       rest of which is the sequence, represented as a string."""
    infile = open(filename)
    info = infile.readline()
    data = infile.read()
    infile.close()
    info = info.replace('\n', '')
    data = data.replace('\n', '')
    info = info.replace('\r', '')	# I hate DOS
    data = data.replace('\r', '')	# I hate DOS
    return [info,data]

def buildTable(rows,cols,val):
    """Build a table with rows rows and cols columns, with each 
       location initialized to val.
       WARNING!  If you try to build a table of lists, there's a
       good chance that they will share memory."""
    table = [[val for col in range(cols)] for row in range(rows)]
    return table

def buildNWTables(seq1,seq2,actionValues):
    """Build the two Needleman-Wunsch tables: One giving the optimal
       alignment value for each pair of subsequences, one giving the
       information on how those alignments are created."""
    len1 = len(seq1)
    len2 = len(seq2)

    # Initialize the two tables
    values = buildTable(len2+1, len1+1, 0)
    directions = buildTable(len2+1, len1+1, 0)
    for row in range(0,len2+1):
        for col in range(0,len1+1):
            directions[row][col] = []

    # Fill in the first row and column of each table
    for col in range(1,len1+1):
        values[0][col] = values[0][col-1] + actionValues['Del']
        directions[0][col] = ['Del']
    for row in range(1,len2+1):
        values[row][0] = values[row-1][0] + actionValues['Ins']
        directions[row][0] = ['Ins']

    # Fill in all the other entries
    for row in range(1,len2+1):
        for col in range(1,len1+1):
            if (seq1[col-1] == seq2[row-1]):
                valMat = values[row-1][col-1] + actionValues['Mat']
                sym = 'Mat'
            else:
                valMat = values[row-1][col-1] + actionValues['Mis']
                sym = 'Mis'
            valIns = values[row-1][col] + actionValues['Ins']
            valDel = values[row][col-1] - actionValues['Del']
            value = max(valMat,valIns,valDel)
            values[row][col] = value
            if (value == valMat):
                directions[row][col].append(sym)
            if (value == valIns):
                directions[row][col].append('Ins')
            if (value == valDel):
                directions[row][col].append('Del')
    # And we're done
    return [values, directions]

def oneAlignmentIterative(directions, row, col):
    """Give one alignment that leads to the match at row/col, using
       directions to figure out the alignment."""
    path = []
    while (row > 0) and (col > 0):
        if 'Del' in directions[row][col]:
            col = col - 1
            path = ['Del'] + path
        elif 'Ins' in directions[row][col]:
            row = row - 1
            path = ['Ins'] + path
        elif 'Mat' in directions[row][col]:
            col = col - 1
            row = row - 1
            path = ['Mat'] + path
        elif 'Mis' in directions[row][col]:
            col = col - 1
            row = row - 1
            path = ['Mis'] + path
        else:
            return ['ERROR']
    return path

def oneAlignment(directions, row, col):
    """Give one alignment that leads to the match at row/col, using
       directions to figure out the alignment."""
    if (row == 0) and (col == 0):
        return []
    elif ('Del') in directions[row][col]:
        return oneAlignment(directions, row, col-1) + ['Del']
    elif ('Ins') in directions[row][col]:
        return oneAlignment(directions, row-1, col) + ['Ins']
    elif ('Mat') in directions[row][col]:
        return oneAlignment(directions, row-1, col-1) + ['Mat']
    elif ('Mis') in directions[row][col]:
        return oneAlignment(directions, row-1, col-1) + ['Mis']
    else:
        return ['Crash']

def allAlignments(directions, row, col):
    """List all the alignments that lead to the match at row/col, using
       directions to figure out the alignment."""
    if (row == 0) and (col == 0):
        return [[]]
    paths = []
    if ('Mat') in directions[row][col]:
        for prev in allAlignments(directions, row-1, col-1):
            paths.append(prev + ['Mat'])
    if ('Mis') in directions[row][col]:
        for prev in allAlignments(directions, row-1, col-1):
            paths.append(prev + ['Mis'])
    if ('Del') in directions[row][col]:
        for prev in allAlignments(directions, row, col-1):
            paths.append(prev + ['Del'])
    if ('Ins') in directions[row][col]:
        for prev in allAlignments(directions, row-1, col):
            paths.append(prev + ['Ins'])
    return paths

def showAlignment(alignment, seq1, seq2):
    """Given an alignment (a sequence of 'Del', 'Ins', 'Mat', and 'Mis') and
       two sequences, show the alignment of those two sequences."""
    pos1 = 0
    pos2 = 0
    aligned1 = []
    aligned2 = []
    notes = []

    # Build the alignment sequences
    for step in alignment:
        if (step == 'Del'):
            aligned1.append(seq1[pos1])
            aligned2.append('-')
            notes.append(' ')
            pos1 = pos1 + 1
        elif (step == 'Ins'):
            aligned1.append('-')
            aligned2.append(seq2[pos2])
            notes.append(' ')
            pos2 = pos2 + 1
        elif ((step == 'Mat') or (step == 'Mis')):
            aligned1.append(seq1[pos1])
            aligned2.append(seq2[pos2])
            if (step == 'Mat'):
                notes.append('*')
            else:
                notes.append(' ')
            pos1 = pos1 + 1
            pos2 = pos2 + 1
        else:
            aligned1.append(' ')
            aligned2.append(' ')
            notes.append('?')
    # Print out the sequences
    print string.join(aligned1, '')
    print string.join(aligned2, '')
    print string.join(notes, '')

def nw1(seq1,seq2,av=actionValuesDefault):
    """Compute and show one alignment of seq1 and seq2 using
       the Needleman-Wunsch algorithm."""
    (values,directions) = buildNWTables(seq1,seq2,av)
    alignment = oneAlignment(directions, len(seq2), len(seq1))
    showAlignment(alignment, seq1, seq2)

def nw2(seq1,seq2,av=actionValuesDefault):
    """Compute and show all alignments of seq1 and seq2 using
       the Needleman-Wunsch algorithm."""
    (values,directions) = buildNWTables(seq1,seq2,av)
    for alignment in allAlignments(directions, len(seq2), len(seq1)):
        showAlignment(alignment, seq1, seq2)
        print ""

def readActionValues():
    """Read the table of action values from the user."""
    av = {}
    av['Del'] = input("Enter the value for deletion: ")
    av['Ins'] = input("Enter the value for insertion: ")
    av['Mat'] = input("Enter the value for a match: ")
    av['Mis'] = input("Enter the value for a mistmatch: ")
    return av

def readMatrix(fname):
    """Read a substitution matrix from the given file."""
    infile = open(fname, 'r')
    matrix = {}
    columns = infile.readline().split()
    lines = infile.readlines()
    for line in lines:
        row = line.split()
        rowlabel = row[0]
        row = row[1:]
        for i in range(0,len(row)):
            key = rowlabel + "." + columns[i]
            matrix[key] = eval(row[i])
    infile.close()
    return matrix

def lookupMatrix(matrix, row, col):
    """Look up a value in the matrix."""
    return matrix[row + '.' + col]

def nwInteractive():
    """Interactively run the Needleman-Wunsch algorithm."""
    av = readActionValues()
    seq1 = raw_input("Enter sequence 1: ")
    seq2 = raw_input("Enter sequence 2: ")
    nw2(seq1, seq2, av)