INT102-Local-Alignment

import numpy as np


def update(a: int, b: int, c: int) -> tuple:
    if a > b and a > c and a > 0:
        return a, "↖"
    if b > a and b > c and b > 0:
        return b, "↑"
    if c > a and c > b and c > 0:
        return c, "←"
    return 0, " "


def mapping(c: str) -> int:
    if c == "A":
        return 0
    elif c == "C":
        return 1
    elif c == "G":
        return 2
    else:
        return 3


def local_alignment(S: list, d: int, X: str, Y: str) -> tuple:
    n, m = len(X), len(Y)
    dp = [[0 for j in range(m + 1)] for i in range(n + 1)]
    flag = [[" " for j in range(m + 1)] for i in range(n + 1)]
    for j in range(m):
        dp[0][j + 1] = max(dp[0][j] + d, 0)
    for i in range(n):
        dp[i + 1][0] = max(dp[i][0] + d, 0)
    for i in range(1, n + 1, 1):
        for j in range(1, m + 1, 1):
            a = mapping(X[i - 1])
            b = mapping(Y[j - 1])
            dp[i][j], flag[i][j] = update(dp[i - 1][j - 1] + S[a][b], dp[i - 1][j] + d, dp[i][j - 1] + d)
    return dp, flag


if __name__ == '__main__':
    X = "GAGT"
    Y = "AGACCT"
    S = [#A   C   G   T
        [ 1, -3, -2, -3],  # A
        [-3,  1, -3, -2],  # C
        [-2, -3,  1, -3],  # G
        [-3, -2, -3,  1]   # T
    ]
    dp, flag = local_alignment(S, -1, X, Y)
    print(np.array(dp))
    print(np.array(flag))