#!/usr/bin/python
import numpy as np
from numpy import array

A, C, G, T = 0, 1, 2, 3
int_to_char = {0:'A', 1:'C', 2:'G', 3:'T'}

indel = -2
score_array = array([[1,-1,-1,-1],[-1,1,-1,-1],[-1,-1,1,-1],[-1,-1,-1,1]])

class alignment(object):
    def __init__(self, sequence1, sequence2):
        self.sequence1 = sequence1
        self.sequence2 = sequence2
        self.D = None

    def global_alignment(self):
        self.D = np.zeros((self.sequence1.size+1, self.sequence2.size+1), dtype=np.int16)
        self.compute_array()
        print self.D
        return self.traceback()

    def get_score(self, i, j):
        return score_array[self.sequence1[i-1], self.sequence2[j-1]]

    def compute_array(self):
        for i in xrange(self.sequence1.size+1):
            self.D[i,0] = i*indel
        for j in xrange(self.sequence2.size+1):
            self.D[0,j] = j*indel
        for i in xrange(1, self.sequence1.size+1):
            for j in xrange(1, self.sequence2.size+1):
                self.D[i,j] = max(self.D[i-1, j-1] + self.get_score(i, j), self.D[i-1, j] + indel, self.D[i, j-1] + indel)
    
    def get_pair(self, i, j):
        n1 = int_to_char[self.sequence1[i-1]] if i>0 else '_'
        n2 = int_to_char[self.sequence2[j-1]] if j>0 else '_'
        return (n1, n2)

    def traceback(self):
        alignment= []
        i = self.sequence1.size
        j = self.sequence2.size
        while i >0 and j>0:
            if self.D[i-1, j-1] + self.get_score(i, j) == self.D[i,j]:
                alignment.append(self.get_pair(i, j))
                i -= 1
                j -= 1
            elif self.D[i-1, j] + indel == self.D[i,j]:
                alignment.append(self.get_pair(i, 0))
                i -= 1
            else:
                alignment.append(self.get_pair(0, j))
                j -= 1
        while i > 0:
            alignment.append(self.get_pair(i, 0))
            i -= 1
        while j > 0:
            alignment.append(self.get_pair(0, j))
            j -= 1
        alignment.reverse()
        return alignment  

def print_seq(pairs):
    top_seq = []
    bottom_seq = []
    for (b, t) in pairs:
        bottom_seq.append(b)
        top_seq.append(t)
    for n in top_seq:
        print n,
    print ' '
    for n in bottom_seq:
        print n,

if __name__ == "__main__":
    seq1 = array([G, C,T,A,A,C,T], dtype=np.int16)
    seq2 = array([C,T,A,A,G,C,T], dtype=np.int16)
    aligner = alignment(seq1, seq2)
    pairs = aligner.global_alignment()
    print_seq(pairs)
