BLOSUM62 = {
'A': {'A': 4, 'C': 0, 'E': -1, 'D': -2, 'G': 0, 'F': -2, 'I': -1, 'H': -2, 'K': -1, 'M': -1, 'L': -1, 'N': -2, 'Q': -1, 'P': -1, 'S': 1, 'R': -1, 'T': 0, 'W': -3, 'V': 0, 'Y': -2},
'C': {'A': 0, 'C': 9, 'E': -4, 'D': -3, 'G': -3, 'F': -2, 'I': -1, 'H': -3, 'K': -3, 'M': -1, 'L': -1, 'N': -3, 'Q': -3, 'P': -3, 'S': -1, 'R': -3, 'T': -1, 'W': -2, 'V': -1, 'Y': -2},
'E': {'A': -1, 'C': -4, 'E': 5, 'D': 2, 'G': -2, 'F': -3, 'I': -3, 'H': 0, 'K': 1, 'M': -2, 'L': -3, 'N': 0, 'Q': 2, 'P': -1, 'S': 0, 'R': 0, 'T': -1, 'W': -3, 'V': -2, 'Y': -2},
'D': {'A': -2, 'C': -3, 'E': 2, 'D': 6, 'G': -1, 'F': -3, 'I': -3, 'H': -1, 'K': -1, 'M': -3, 'L': -4, 'N': 1, 'Q': 0, 'P': -1, 'S': 0, 'R': -2, 'T': -1, 'W': -4, 'V': -3, 'Y': -3},
'G': {'A': 0, 'C': -3, 'E': -2, 'D': -1, 'G': 6, 'F': -3, 'I': -4, 'H': -2, 'K': -2, 'M': -3, 'L': -4, 'N': 0, 'Q': -2, 'P': -2, 'S': 0, 'R': -2, 'T': -2, 'W': -2, 'V': -3, 'Y': -3},
'F': {'A': -2, 'C': -2, 'E': -3, 'D': -3, 'G': -3, 'F': 6, 'I': 0, 'H': -1, 'K': -3, 'M': 0, 'L': 0, 'N': -3, 'Q': -3, 'P': -4, 'S': -2, 'R': -3, 'T': -2, 'W': 1, 'V': -1, 'Y': 3},
'I': {'A': -1, 'C': -1, 'E': -3, 'D': -3, 'G': -4, 'F': 0, 'I': 4, 'H': -3, 'K': -3, 'M': 1, 'L': 2, 'N': -3, 'Q': -3, 'P': -3, 'S': -2, 'R': -3, 'T': -1, 'W': -3, 'V': 3, 'Y': -1},
'H': {'A': -2, 'C': -3, 'E': 0, 'D': -1, 'G': -2, 'F': -1, 'I': -3, 'H': 8, 'K': -1, 'M': -2, 'L': -3, 'N': 1, 'Q': 0, 'P': -2, 'S': -1, 'R': 0, 'T': -2, 'W': -2, 'V': -3, 'Y': 2},
'K': {'A': -1, 'C': -3, 'E': 1, 'D': -1, 'G': -2, 'F': -3, 'I': -3, 'H': -1, 'K': 5, 'M': -1, 'L': -2, 'N': 0, 'Q': 1, 'P': -1, 'S': 0, 'R': 2, 'T': -1, 'W': -3, 'V': -2, 'Y': -2},
'M': {'A': -1, 'C': -1, 'E': -2, 'D': -3, 'G': -3, 'F': 0, 'I': 1, 'H': -2, 'K': -1, 'M': 5, 'L': 2, 'N': -2, 'Q': 0, 'P': -2, 'S': -1, 'R': -1, 'T': -1, 'W': -1, 'V': 1, 'Y': -1},
'L': {'A': -1, 'C': -1, 'E': -3, 'D': -4, 'G': -4, 'F': 0, 'I': 2, 'H': -3, 'K': -2, 'M': 2, 'L': 4, 'N': -3, 'Q': -2, 'P': -3, 'S': -2, 'R': -2, 'T': -1, 'W': -2, 'V': 1, 'Y': -1},
'N': {'A': -2, 'C': -3, 'E': 0, 'D': 1, 'G': 0, 'F': -3, 'I': -3, 'H': 1, 'K': 0, 'M': -2, 'L': -3, 'N': 6, 'Q': 0, 'P': -2, 'S': 1, 'R': 0, 'T': 0, 'W': -4, 'V': -3, 'Y': -2},
'Q': {'A': -1, 'C': -3, 'E': 2, 'D': 0, 'G': -2, 'F': -3, 'I': -3, 'H': 0, 'K': 1, 'M': 0, 'L': -2, 'N': 0, 'Q': 5, 'P': -1, 'S': 0, 'R': 1, 'T': -1, 'W': -2, 'V': -2, 'Y': -1},
'P': {'A': -1, 'C': -3, 'E': -1, 'D': -1, 'G': -2, 'F': -4, 'I': -3, 'H': -2, 'K': -1, 'M': -2, 'L': -3, 'N': -2, 'Q': -1, 'P': 7, 'S': -1, 'R': -2, 'T': -1, 'W': -4, 'V': -2, 'Y': -3},
'S': {'A': 1, 'C': -1, 'E': 0, 'D': 0, 'G': 0, 'F': -2, 'I': -2, 'H': -1, 'K': 0, 'M': -1, 'L': -2, 'N': 1, 'Q': 0, 'P': -1, 'S': 4, 'R': -1, 'T': 1, 'W': -3, 'V': -2, 'Y': -2},
'R': {'A': -1, 'C': -3, 'E': 0, 'D': -2, 'G': -2, 'F': -3, 'I': -3, 'H': 0, 'K': 2, 'M': -1, 'L': -2, 'N': 0, 'Q': 1, 'P': -2, 'S': -1, 'R': 5, 'T': -1, 'W': -3, 'V': -3, 'Y': -2},
'T': {'A': 0, 'C': -1, 'E': -1, 'D': -1, 'G': -2, 'F': -2, 'I': -1, 'H': -2, 'K': -1, 'M': -1, 'L': -1, 'N': 0, 'Q': -1, 'P': -1, 'S': 1, 'R': -1, 'T': 5, 'W': -2, 'V': 0, 'Y': -2},
'W': {'A': -3, 'C': -2, 'E': -3, 'D': -4, 'G': -2, 'F': 1, 'I': -3, 'H': -2, 'K': -3, 'M': -1, 'L': -2, 'N': -4, 'Q': -2, 'P': -4, 'S': -3, 'R': -3, 'T': -2, 'W': 11, 'V': -3, 'Y': 2},
'V': {'A': 0, 'C': -1, 'E': -2, 'D': -3, 'G': -3, 'F': -1, 'I': 3, 'H': -3, 'K': -2, 'M': 1, 'L': 1, 'N': -3, 'Q': -2, 'P': -2, 'S': -2, 'R': -3, 'T': 0, 'W': -3, 'V': 4, 'Y': -1},
'Y': {'A': -2, 'C': -2, 'E': -2, 'D': -3, 'G': -3, 'F': 3, 'I': -1, 'H': 2, 'K': -2, 'M': -1, 'L': -1, 'N': -2, 'Q': -1, 'P': -3, 'S': -2, 'R': -2, 'T': -2, 'W': 2, 'V': -1, 'Y': 7}}


def output_lcs(bt_lower, bt_middle, bt_upper, v, w, i, j):

    print("BT Lower")
    for s in bt_lower:
        print(s)
    print("BT middle: ")
    for s in bt_middle:
        print(s)
    print("upper")
    for s in bt_upper:
        print(s)

    w_trace = []
    v_trace = []
    cur_m = 'middle'
    while True:
        print("i j", i, j)
        if i == 0 or j == 0:
            break
        if cur_m == 'middle':
            if bt_middle[i][j] == '↓':
                v_trace.append(v[i-1])
                w_trace.append('-')
                #i -= 1
                cur_m = 'lower'
            elif bt_middle[i][j] == '→':
                w_trace.append(w[j-1])
                v_trace.append('-')
                #j -= 1
                cur_m = 'upper'
            else:
                v_trace.append(v[i-1])
                w_trace.append(w[j-1])
                i -= 1
                j -= 1
        elif cur_m == 'upper':
            if bt_upper[i][j] == '→':
                w_trace.append(w[j-1])
                v_trace.append('-')
                j -= 1
            else:
                v_trace.append(v[i-1])
                w_trace.append(w[j-1])
                #i -= 1
                j -= 1
                cur_m == 'middle'
        elif cur_m == 'lower':
            if bt_lower[i][j] == '↓':
                v_trace.append(v[i-1])
                w_trace.append('-')
                i -= 1
            else:
                v_trace.append(v[i-1])
                w_trace.append(w[j-1])
                i -= 1
                cur_m = 'middle'


    # Start of the strings
    if i == 0 and j != 0:
        w_trace.append(w[j-1])
        v_trace.append('-')
    elif j == 0 and i != 0:
        v_trace.append(v[i-1])
        w_trace.append('-')

    v_res = ''.join(reversed(v_trace))
    w_res = ''.join(reversed(w_trace))
    print("V: ", v_res)
    print("W: ", w_res)
    return v_res, w_res


def align_affine_gap(v, w):
    sigma = 11
    eps = 1

    m_lower = [[0] * (len(w)+1) for _ in range(len(v)+1)]
    m_middle = [[0] * (len(w)+1) for _ in range(len(v)+1)]
    m_upper = [[0] * (len(w)+1) for _ in range(len(v)+1)]

    bt_lower = [['*'] * (len(w)+1) for _ in range(len(v)+1)]
    bt_upper = [['*'] * (len(w)+1) for _ in range(len(v)+1)]
    bt_middle = [['*'] * (len(w)+1) for _ in range(len(v)+1)]

    # Initialize matrices' boundaries
    for i in range(1, len(v)+1):
        if i == 1:
            m_lower[i][0] = m_lower[i-1][0] - sigma
            m_middle[i][0] = m_middle[i-1][0] - sigma
        else:
            m_lower[i][0] = m_lower[i-1][0] - eps
            m_middle[i][0] = m_middle[i-1][0] - eps
    for j in range(1, len(w)+1):
        if j == 1:
            m_upper[0][j] = m_upper[0][j-1] - sigma
            m_middle[0][j] = m_middle[0][j-1] - sigma
        else:
            m_upper[0][j] = m_upper[0][j-1] - eps
            m_middle[0][j] = m_middle[0][j-1] - eps

    for i in range(1, len(v)+1):
        for j in range(1, len(w)+1):
            match_score = BLOSUM62[v[i-1]][w[j-1]]

            m_lower[i][j] = max(m_lower[i-1][j]-eps, m_middle[i-1][j]-sigma)
            m_upper[i][j] = max(m_upper[i][j-1]-eps, m_middle[i][j-1]-sigma)
            m_middle[i][j] = max(m_lower[i][j], m_upper[i][j], m_middle[i-1][j-1] + match_score)

            # Backtrack path
            if m_lower[i][j] == m_lower[i-1][j]-eps:
                bt_lower[i][j] = '↓'
            else:
                bt_lower[i][j] = '~'
            if m_upper[i][j] == m_upper[i][j-1]-eps:
                bt_upper[i][j] = '→'
            else:
                bt_upper[i][j] = '~'
            if m_middle[i][j] == m_lower[i][j]:
                bt_middle[i][j] = '↓'
            elif m_middle[i][j] == m_upper[i][j]:
                bt_middle[i][j] = '→'
            else:
                bt_middle[i][j] = '↘'

    print("Upper: ")
    for i in m_upper:
        print(i)
    print("Middle")
    for i in m_middle:
        print(i)
    print("Lower")
    for i in m_lower:
        print(i)

    v_res, w_res = output_lcs(bt_lower, bt_middle, bt_upper, v, w, len(v), len(w))

    return v_res, w_res, m_middle[-1][-1]


if __name__ == "__main__":
    # (1) Alignment with affine gap penalties problem
    w = 'PRTEINS'
    v = 'PRTWPSEIN'
    score, v_al, w_al = align_affine_gap(v, w)
    print(score)
    print(v_al)
    print(w_al)
