import matplotlib.pyplot as plt
import numpy as np

#定義全局變量列表finalList存儲最后回溯的路徑 finalOrder1,finalOrder2存儲最后的序列 finalRoad用于存儲方向路徑用于最后畫圖
def createList():
  global finalList
  global finalOrder1
  global finalOrder2
  global finalRoad
  finalList = []
  finalOrder1 = []
  finalOrder2 = []
  finalRoad = []


#創(chuàng)建A G C T 對應(yīng)的鍵值對，方便查找計分矩陣中對應(yīng)的得分
def createDic():
  dic = {'A':0,'G':1,'C':2,'T':3}
  return dic
#構(gòu)建計分矩陣
# A G C T
def createGrade():
  grade = np.matrix([[10,-1,-3,-4],
            [-1,7,-5,-3],
            [-3,-5,9,0],
            [-4,-3,0,8]])
  return grade

#計算兩個字符的相似度得分函數(shù)
def getGrade(a,b):
  dic = createDic() # 堿基字典 方便查找計分矩陣
  grade = createGrade() # 打分矩陣grade
  return grade[dic[a],dic[b]]

#構(gòu)建得分矩陣函數(shù) 參數(shù)為要比較序列、自定義的gap值
def createMark(s,t,gap):
  a = len(s)             #獲取序列長度a,b
  b = len(t)
  mark = np.zeros((a+1,b+1))     #初始化全零得分矩陣
  direction = np.zeros((a+1,b+1,3)) #direction矩陣用來存儲得分矩陣中得分來自的方向  第一個表示左方 第二個表示左上 第三個表示上方 1表示能往哪個方向去
                    #由于得分可能會來自多個方向，所以使用三維矩陣存儲

  direction[0][0] = -1        #確定回溯時的結(jié)束條件 即能夠走到方向矩陣的值為-1
  mark[0,:] = np.fromfunction(lambda x, y: gap * (x + y), (1, b + 1), dtype=int)   #根據(jù)gap值將得分矩陣第一行計算出
  mark[:,0] = np.fromfunction(lambda x, y: gap * (x + y), (1, a + 1), dtype=int)   #根據(jù)gap值將得分矩陣第一列計算出
  for i in range(1,b+1):
    direction[0,i,0] = 1
  for i in range(1, a + 1):
    direction[i, 0, 2] = 1

  for i in range(1,a+1):
    for j in range(1,b+1):
      threeMark = [mark[i][j-1],mark[i-1][j-1],mark[i-1][j]]     #threeMark表示現(xiàn)在所要計算得分的位置的左邊 左上 上邊的得分
      threeGrade = [gap,getGrade(s[i-1],t[j-1]),gap]         #threeGrade表示經(jīng)過需要計算得左邊 左上 上邊的空位以及相似度得分
      finalGrade = np.add(threeMark,threeGrade)           #finalGrade表示最終來自三個方向上的得分
      mark[i][j] = max(finalGrade)                  #選取三個方向上的最大得分存入得分矩陣
      #可能該位置的得分可以由多個方向得來，所以進(jìn)行判斷并循環(huán)賦值
      for k in range(0,len([y for y,x in enumerate(finalGrade) if x == max(finalGrade)])):
        directionList = [y for y,x in enumerate(finalGrade) if x == max(finalGrade)]
        direction[i][j][directionList[k]] = 1
  return mark,direction

#回溯函數(shù) 參數(shù)分別為 得分矩陣 方向矩陣 現(xiàn)在所處得分矩陣的位置 以及兩個序列
def remount(mark,direction,i,j,s,t):
  if direction[i][j][0] == 1 :
    if direction[i][j-1][0] == -1:      #如果該位置指向左邊 先判斷其左邊是否是零點(diǎn)
      finalList.append(0)          #如果是 將該路徑存入路徑列表
      finalList.reverse()          #將列表反過來得到從零點(diǎn)開始的路徑
      index1 = 0              #記錄現(xiàn)在所匹配序列s的位置 因為兩個字符串可能是不一樣長的
      index2 = 0              #記錄現(xiàn)在所匹配序列t的位置
      for k in finalList:
        if k == 0 :
          finalOrder1.append("-")
          finalOrder2.append(t[index2])
          index2 += 1
        if k == 1 :
          finalOrder1.append(s[index1])
          finalOrder2.append(t[index2])
          index1 += 1
          index2 += 1
        if k == 2 :
          finalOrder1.append(s[index1])
          finalOrder2.append("-")
          index1 += 1
      finalList.reverse() # 將原來反轉(zhuǎn)的路徑再返回來
      finalRoad.append(np.array(finalList)) # 將此次的路徑添加到最終路徑記錄用于最后畫圖
      finalList.pop()            #輸出后將當(dāng)前方向彈出 并回溯
      return
    else :
      finalList.append(0)          #如果不是零點(diǎn) 則將該路徑加入路徑矩陣，繼續(xù)往下走
      remount(mark,direction,i,j-1,s,t)
      finalList.pop()            #該方向走完后將這個方向彈出 繼續(xù)下一輪判斷 下面兩個大的判斷同理
  if direction[i][j][1] == 1 :
    if direction[i-1][j-1][0] == -1:

      finalList.append(1)
      finalList.reverse()          # 將列表反過來得到從零點(diǎn)開始的路徑
      index1 = 0              # 記錄現(xiàn)在所匹配序列s的位置 因為兩個字符串可能是不一樣長的
      index2 = 0              # 記錄現(xiàn)在所匹配序列t的位置
      for k in finalList:
        if k == 0 :
          finalOrder1.append("-")
          finalOrder2.append(t[index2])
          index2 += 1
        if k == 1 :
          finalOrder1.append(s[index1])
          finalOrder2.append(t[index2])
          index1 += 1
          index2 += 1
        if k == 2 :
          finalOrder1.append(s[index1])
          finalOrder2.append("-")
          index1 += 1
      finalList.reverse() # 將原來反轉(zhuǎn)的路徑再返回來
      finalRoad.append(np.array(finalList)) # 將此次的路徑添加到最終路徑記錄用于最后畫圖
      finalList.pop()
      return
    else :
      finalList.append(1)
      remount(mark,direction,i-1,j-1,s,t)
      finalList.pop()
  if direction[i][j][2] == 1 :
    if direction[i-1][j][0] == -1:
      finalList.append(2)
      finalList.reverse()           # 將列表反過來得到從零點(diǎn)開始的路徑
      index1 = 0                # 記錄現(xiàn)在所匹配序列s的位置 因為兩個字符串可能是不一樣長的
      index2 = 0                # 記錄現(xiàn)在所匹配序列t的位置
      for k in finalList:
        if k == 0 :
          finalOrder1.append("-")
          finalOrder2.append(t[index2])
          index2 += 1
        if k == 1 :
          finalOrder1.append(s[index1])
          finalOrder2.append(t[index2])
          index1 += 1
          index2 += 1
        if k == 2 :
          finalOrder1.append(s[index1])
          finalOrder2.append("-")
          index1 += 1
      finalList.reverse()          # 將原來反轉(zhuǎn)的路徑再返回來
      finalRoad.append(np.array(finalList)) # 將此次的路徑添加到最終路徑記錄用于最后畫圖
      finalList.pop()
      return
    else :
      finalList.append(2)
      remount(mark,direction,i-1,j,s,t)
      finalList.pop()

#畫箭頭函數(shù)
def arrow(ax,sX,sY,aX,aY):
  ax.arrow(sX,sY,aX,aY,length_includes_head=True, head_width=0.15, head_length=0.25, fc='w', ec='b')

#畫圖函數(shù)
def drawArrow(mark, direction, a, b, s, t):
  #a是s的長度為4  b是t的長度為6
  fig = plt.figure()
  ax = fig.add_subplot(111)
  val_ls = range(a+2)
  scale_ls = range(b+2)
  index_ls = []
  index_lsy = []
  for i in range(a):
    if i == 0:
      index_lsy.append('#')
    index_lsy.append(s[a-i-1])
  index_lsy.append('0')
  for i in range(b):
    if i == 0:
      index_ls.append('#')
      index_ls.append('0')
    index_ls.append(t[i])
  plt.xticks(scale_ls, index_ls)      #設(shè)置坐標(biāo)字
  plt.yticks(val_ls, index_lsy)
  for k in range(1,a+2):
    y = [k for i in range(0,b+1)]
    x = [x for x in range(1,b+2)]
    ax.scatter(x, y, c='y')
  for i in range(1,a+2):
    for j in range(1,b+2):
      ax.text(j,a+2-i,int(mark[i-1][j-1]))
  lX = b+1
  lY = 1
  for n in range(0,len(finalRoad)):
    for m in (finalRoad[n]):
      if m == 0:
        arrow(ax,lX,lY,-1,0)
        lX = lX - 1
      elif m == 1:
        arrow(ax,lX,lY,-1,1)
        lX = lX - 1
        lY = lY + 1
      elif m == 2:
        arrow(ax, lX, lY, 0, 1)
        lY = lY + 1
    lX = b + 1
    lY = 1
  ax.set_xlim(0, b + 2) # 設(shè)置圖形的范圍，默認(rèn)為[0,1]
  ax.set_ylim(0, a + 2) # 設(shè)置圖形的范圍，默認(rèn)為[0,1]
  ax.set_aspect('equal') # x軸和y軸等比例
  plt.show()
  plt.tight_layout()
if __name__ == '__main__':
  createList()
  print("Please enter gap:")
  gap = int(input())       #獲取gap值 轉(zhuǎn)換為整型  tip：剛開始就是因為這里沒有進(jìn)行類型導(dǎo)致后面的計算部分報錯
  print("Please enter sequence 1:")
  s = input()           #獲取用戶輸入的第一條序列
  print("Please enter sequence 2:")
  t = input()           #獲取用戶輸入的第二條序列
  a = len(s)           #獲取s的長度
  b = len(t)           #獲取t的長度
  mark,direction = createMark(s,t,gap)
  print("The scoring matrix is as follows：")     #輸出得分矩陣
  print(mark)
  remount(mark,direction,a,b,s,t) #調(diào)用回溯函數(shù)
  c = a if a > b else b     #判斷有多少種比對結(jié)果得到最終比對序列的長度
  total = int(len(finalOrder1)/c)
  for i in range(1,total+1):   #循環(huán)輸出比對結(jié)果
    k = str(i)
    print("Sequence alignment results "+k+" is：")
    print(finalOrder1[(i-1)*c:i*c])
    print(finalOrder2[(i-1)*c:i*c])
  drawArrow(mark, direction, a, b, s, t)

import matplotlib.pyplot as plt
import numpy as np
import operator
#在局部比對中 回溯結(jié)束的條件是方向矩陣中該位置的值全為0
#定義全局變量列表finalList存儲最后回溯的路徑 finalOrder1,finalOrder2存儲最后的序列
def createList():
  global finalList
  global finalOrder1
  global finalOrder2
  global finalRoad
  finalList = []
  finalOrder1 = []
  finalOrder2 = []
  finalRoad = []


#創(chuàng)建A G C T 對應(yīng)的鍵值對，方便查找計分矩陣中對應(yīng)的得分
def createDic():
  dic = {'A':0,'G':1,'C':2,'T':3}
  return dic
#構(gòu)建計分矩陣
# A G C T
def createGrade():
  grade = np.matrix([[10,-1,-3,-4],
            [-1,7,-5,-3],
            [-3,-5,9,0],
            [-4,-3,0,8]])
  return grade

#計算兩個字符的相似度得分函數(shù)
def getGrade(a,b):
  dic = createDic() # 堿基字典 方便查找計分矩陣
  grade = createGrade() # 打分矩陣grade
  return grade[dic[a],dic[b]]

#構(gòu)建得分矩陣函數(shù) 參數(shù)為要比較序列、自定義的gap值
def createMark(s,t,gap):
  a = len(s)             #獲取序列長度a,b
  b = len(t)
  mark = np.zeros((a+1,b+1))     #初始化全零得分矩陣
  direction = np.zeros((a+1,b+1,3)) #direction矩陣用來存儲得分矩陣中得分來自的方向  第一個表示左方 第二個表示左上 第三個表示上方 1表示能往哪個方向去
                    #由于得分可能會來自多個方向，所以使用三維矩陣存
  for i in range(1,a+1):
    for j in range(1,b+1):
      threeMark = [mark[i][j-1],mark[i-1][j-1],mark[i-1][j]]     #threeMark表示現(xiàn)在所要計算得分的位置的左邊 左上 上邊的得分
      threeGrade = [gap,getGrade(s[i-1],t[j-1]),gap]         #threeGrade表示經(jīng)過需要計算得左邊 左上 上邊的空位以及相似度得分
      finalGrade = np.add(threeMark,threeGrade)           #finalGrade表示最終來自三個方向上的得分
      if max(finalGrade) >= 0:                  #如果該最大值是大于0的則 選取三個方向上的最大得分存入得分矩陣 否則不對矩陣進(jìn)行修改
        mark[i][j] = max(finalGrade)
        for k in range(0,len([y for y,x in enumerate(finalGrade) if x == max(finalGrade)])): #可能該位置的得分可以由多個方向得來，所以進(jìn)行判斷并循環(huán)賦值
          directionList = [y for y,x in enumerate(finalGrade) if x == max(finalGrade)]
          direction[i][j][directionList[k]] = 1
  return mark,direction

#回溯函數(shù) 參數(shù)分別為 得分矩陣 方向矩陣 現(xiàn)在所處得分矩陣的位置 以及兩個序列
def remount(mark,direction,i,j,s,t):
  if direction[i][j][0] == 1 :
    if all(direction[i][j-1] == [0,0,0]):      #如果該位置指向左邊 先判斷其左邊是否是零點(diǎn)
      finalList.append(0)          #如果是 將該路徑存入路徑列表
      finalList.reverse()          #將列表反過來得到從零點(diǎn)開始的路徑
      index1 = i              #記錄現(xiàn)在所匹配序列s的位置 因為兩個字符串可能是不一樣長的
      index2 = j-1              #記錄現(xiàn)在所匹配序列t的位置
      for k in finalList:
        if k == 0 :
          finalOrder1.append("-")
          finalOrder2.append(t[index2])
          index2 += 1
        if k == 1 :
          finalOrder1.append(s[index1])
          finalOrder2.append(t[index2])
          index1 += 1
          index2 += 1
        if k == 2 :
          finalOrder1.append(s[index1])
          finalOrder2.append("-")
          index1 += 1
      finalList.reverse()
      finalRoad.append(np.array(finalList)) # 將此次的路徑添加到最終路徑記錄用于最后畫圖
      finalList.pop()            #輸出后將當(dāng)前方向彈出 并回溯
      return
    else :
      finalList.append(0)          #如果不是零點(diǎn) 則將該路徑加入路徑矩陣，繼續(xù)往下走
      remount(mark,direction,i,j-1,s,t)
      finalList.pop()            #該方向走完后將這個方向彈出 繼續(xù)下一輪判斷 下面兩個大的判斷同理
  if direction[i][j][1] == 1 :
    if all(direction[i-1][j-1] == [0,0,0]):
      finalList.append(1)
      finalList.reverse()           # 將列表反過來得到從零點(diǎn)開始的路徑
      index1 = i-1              # 記錄現(xiàn)在所匹配序列s的位置 因為兩個字符串可能是不一樣長的
      index2 = j-1              # 記錄現(xiàn)在所匹配序列t的位置
      for k in finalList:
        if k == 0 :
          finalOrder1.append("-")
          finalOrder2.append(t[index2])
          index2 += 1
        if k == 1 :
          finalOrder1.append(s[index1])
          finalOrder2.append(t[index2])
          index1 += 1
          index2 += 1
        if k == 2 :
          finalOrder1.append(s[index1])
          finalOrder2.append("-")
          index1 += 1
      finalList.reverse()
      finalRoad.append(np.array(finalList)) # 將此次的路徑添加到最終路徑記錄用于最后畫圖
      finalList.pop()
      return
    else :
      finalList.append(1)
      remount(mark,direction,i-1,j-1,s,t)
      finalList.pop()
  if direction[i][j][2] == 1 :
    if all(direction[i-1][j] == [0,0,0]):
      finalList.append(2)
      finalList.reverse()           # 將列表反過來得到從零點(diǎn)開始的路徑
      index1 = i-1                # 記錄現(xiàn)在所匹配序列s的位置 因為兩個字符串可能是不一樣長的
      index2 = j                # 記錄現(xiàn)在所匹配序列t的位置
      for k in finalList:
        if k == 0 :
          finalOrder1.append("-")
          finalOrder2.append(t[index2])
          index2 += 1
        if k == 1 :
          finalOrder1.append(s[index1])
          finalOrder2.append(t[index2])
          index1 += 1
          index2 += 1
        if k == 2 :
          finalOrder1.append(s[index1])
          finalOrder2.append("-")
          index1 += 1
      finalList.reverse()
      finalRoad.append(np.array(finalList)) # 將此次的路徑添加到最終路徑記錄用于最后畫圖
      finalList.pop()
      return
    else :
      finalList.append(2)
      remount(mark,direction,i-1,j,s,t)
      finalList.pop()
#畫箭頭函數(shù)
def arrow(ax,sX,sY,aX,aY):
  ax.arrow(sX,sY,aX,aY,length_includes_head=True, head_width=0.15, head_length=0.25, fc='w', ec='b')

#畫圖函數(shù)
def drawArrow(mark, direction, a, b, s, t,mx,my):
  #a是s的長度為4  b是t的長度為6
  fig = plt.figure()
  ax = fig.add_subplot(111)
  val_ls = range(a+2)
  scale_ls = range(b+2)
  index_ls = []
  index_lsy = []
  for i in range(a):
    if i == 0:
      index_lsy.append('#')
    index_lsy.append(s[a-i-1])
  index_lsy.append('0')
  for i in range(b):
    if i == 0:
      index_ls.append('#')
      index_ls.append('0')
    index_ls.append(t[i])
  plt.xticks(scale_ls, index_ls)      #設(shè)置坐標(biāo)字
  plt.yticks(val_ls, index_lsy)
  for k in range(1,a+2):
    y = [k for i in range(0,b+1)]
    x = [x for x in range(1,b+2)]
    ax.scatter(x, y, c='y')
  for i in range(1,a+2):
    for j in range(1,b+2):
      ax.text(j,a+2-i,int(mark[i-1][j-1]))
  lX = my + 1
  lY = a - mx + 1
  for n in range(0,len(finalRoad)):
    for m in (finalRoad[n]):
      if m == 0:
        arrow(ax,lX,lY,-1,0)
        lX = lX - 1
      elif m == 1:
        arrow(ax,lX,lY,-1,1)
        lX = lX - 1
        lY = lY + 1
      elif m == 2:
        arrow(ax, lX, lY, 0, 1)
        lY = lY + 1
    lX = b + 1
    lY = 1
  ax.set_xlim(0, b + 2) # 設(shè)置圖形的范圍，默認(rèn)為[0,1]
  ax.set_ylim(0, a + 2) # 設(shè)置圖形的范圍，默認(rèn)為[0,1]
  ax.set_aspect('equal') # x軸和y軸等比例
  plt.show()
  plt.tight_layout()

if __name__ == '__main__':
  createList()
  print("Please enter gap:")
  gap = int(input())       #獲取gap值 轉(zhuǎn)換為整型  tip：剛開始就是因為這里沒有進(jìn)行類型導(dǎo)致后面的計算部分報錯
  print("Please enter sequence 1:")
  s = input()           #獲取用戶輸入的第一條序列
  print("Please enter sequence 2:")
  t = input()           #獲取用戶輸入的第二條序列
  a = len(s)           #獲取s的長度
  b = len(t)           #獲取t的長度
  mark,direction = createMark(s,t,gap)
  print("The scoring matrix is as follows：")     #輸出得分矩陣
  print(mark)
  maxDirection = np.argmax(mark) #獲取最大值的位置
  i = int(maxDirection/(b+1))
  j = int(maxDirection - i*(b+1))
  remount(mark,direction,i,j,s,t) #調(diào)用回溯函數(shù)
  print(finalOrder1)
  print(finalOrder2)
  drawArrow(mark, direction, a, b, s, t, i, j)