from Bio import SeqIO

#--------------------------------------------------------------------------------------
# read PAM250 alignment scoring matrix
# return matrix and symbols lookup dict

def read_pam_table():
    f = open('helpers/PAM250.txt')
    lines = f.readlines()
    
    # 'A': 0, 'C': 1, 'D': 2...
    symbols_dict = {}
    symbols = lines[0].rstrip().split(' ')
    i = 0
    for symbol in symbols:
        symbols_dict[symbol] = i
        i = i + 1

    n = len(symbols_dict)
    B = [[0] * n for i in range(n)]
    for i in range(n):
        scores = lines[i+1].rstrip().split(' ')
        for j in range(n):
            B[i][j] = int(scores[j])
    return B, symbols_dict

#--------------------------------------------------------------------------------------
# initialize and return dynamic programming table M

def initialize():
    # create a matrix M of sizes w x h
    M = [[0] * w for i in range(h)]
    # same for directions used to trace-back an optimal alignment
    directions = [[0] * w for i in range(h)]
    # fill entries in the first column and row (punishments for introducing an indel)
    for i in range(h):
        M[i][0] = -5*i
    for j in range(w):
        M[0][j] = -5*j
    return M, directions

#--------------------------------------------------------------------------------------
# fill dynamic programming table M based on score function
# keep track of best alignment score

def compute_entries():
    global max_score, max_score_i, max_score_j
    for i in range(1, h):
        for j in range(1, w):
            [s_v, s_h, s_d] = compute_scores(i, j)
            scores = [s_v, s_h, s_d, 0]
            best_score_val = max(scores)
            best_score_dir = scores.index(best_score_val)
            directions[i][j] = best_score_dir
            M[i][j] = best_score_val
            if (best_score_val > max_score):
                max_score = best_score_val
                max_score_i = i
                max_score_j = j

#--------------------------------------------------------------------------------------
# compute scores for all different moves (vertical, horizontal, diagonal)

def compute_scores(i, j):
    s_v = M[i-1][j] + get_score(s[i], '-')
    s_h = M[i][j-1] + get_score('-', t[j])
    s_d = M[i-1][j-1] + get_score(s[i], t[j])
    # s_v = M[i-1][j] + get_score('-', t[j])
    # s_h = M[i][j-1] + get_score(s[i], '-')
    # s_d = M[i-1][j-1] + get_score(s[i], t[j])

    return [s_v, s_h, s_d]

#--------------------------------------------------------------------------------------
# compare symbols and return score based on scoring function

def get_score(a, b):
    if (a == '-' or b == '-'):
        return -5
    
    i = symbols_dict[a]
    j = symbols_dict[b]
    return B[i][j]

#--------------------------------------------------------------------------------------
# reconstruct the optimal alignment of the strings with highest score

def traceback():
    s_o = []
    t_o = []
    i = max_score_i
    j = max_score_j
    
    while True:
        dir = directions[i][j]
        if (dir == 0):
            s_o.insert(0, s[i])
            t_o.insert(0, '-')
            i = i-1
        elif (dir == 1):
            s_o.insert(0, '-')
            t_o.insert(0, t[j])
            j = j-1
        elif (dir == 2):
            s_o.insert(0, s[i])
            t_o.insert(0, t[j])
            i = i-1
            j = j-1
        else:
            break

        if (i == 0 or j == 0):
            break

    return [s_o, t_o]

#--------------------------------------------------------------------------------------

solution = open('solutions/24-solution.txt', 'w')

records = []
for record in SeqIO.parse("datasets/24-dataset.txt", "fasta"):
    records.append(str(record.seq))

s = '-' + records[0]
t = '-' + records[1]

w = len(t)
h = len(s)

# PAM250 scoring matrix
B, symbols_dict = read_pam_table()

# prepare variables for finding max alignment score
max_score = 0
max_score_i = 0
max_score_j = 0

# Needleman-Wunsch algorithm for finding local alignment of substrings
M, directions = initialize()
compute_entries()

# find substrings that produce the maximum alignment score
s_o, t_o = traceback()

solution.write(str(max_score) + '\n')

for c in s_o:
    if (c != '-'):
        solution.write(c)
solution.write('\n')
for c in t_o:
    if (c != '-'):
        solution.write(c)

solution.close()