import numpy as np
import numba
from numba import jit

@jit
def Needleman(a,b):
    medzera_open=-3
    medzera=-1
    pair={
    "AA":1,"AC":-1,"AG":-1,"AT":-1,
    "CA":-1,"CC":1,"CG":-1,"CT":-1,
    "GA":-1,"GC":-1,"GG":1,"GT":-1,
    "TA":-1,"TC":-1,"TG":-1,"TT":1,
    }
    A=np.zeros((len(a)+1,len(b)+1),dtype='int32')
    B=np.zeros((len(a)+1,len(b)+1),dtype='int32')
    B[0,0]=-1
    for i in range(len(A)):
        A[i][0]=i*medzera+medzera_open
        B[i][0]=2
    for i in range(len(A[0])):
        A[0][i]=i*medzera+medzera_open
        B[0][i]=3
    for i in range(1,len(A)):
        for j in range(1,len(A[0])):
            akt_medzera_up=medzera+ (medzera_open if  B[i-1,j]!=2 else 0)
            akt_medzera_left=medzera+ (medzera_open if  B[i,j-1]!=3 else 0)
            A[i,j]=max(A[i-1,j-1]+pair[a[i-1]+b[j-1]],A[i-1,j]+akt_medzera_up,A[i,j-1]+akt_medzera_left)
            B[i,j]=np.argmax(np.asarray((A[i-1,j-1]+pair[a[i-1]+b[j-1]],A[i-1,j]+akt_medzera_up,A[i,j-1]+akt_medzera_left)))+1
    r=B.shape[0]-1
    col=B.shape[1]-1
    res=np.ones((len(a)+len(b)+2),dtype='int32')*-1
    pos=res.shape[0]-1
    while r>0 or col>0:
        res[pos]=B[r,col]
        pos-=1
        if B[r,col]==1:
            r-=1
            col-=1
        if B[r,col]==2:
            r-=1
        if B[r,col]==3:
            col-=1
    return res[pos+1:]

def transform(seq,multiplier=1):
    ret=np.zeros(len(seq),dtype='int32').reshape((len(seq)))
    dc={'A':1,'T':2,'G':3,'C':4}
    for i in range(len(seq)):
        ret[i]=dc[seq[i]]*multiplier
    return ret
        
with open('data_fixed') as vstup:
        count=0
        with open(f'AllignedFixed/out.txt','w') as vystup:
            for row in vstup:
                print(count)
                count+=1
                akt=row.split()
                orig=akt[0]
                a=akt[1]
                b=akt[2]
                alignment=list(Needleman(orig,a))
                seqs=list(transform(orig)+transform(a,multiplier=10))
                vystup.write(str(seqs)+' ')
                vystup.write(str(alignment)+'\n')
                alignment=list(Needleman(orig,b))
                seqs=list(transform(orig)+transform(b,multiplier=10))
                vystup.write(str(seqs)+' ')
                vystup.write(str(alignment)+'\n')
            
           
