"""
Function of this code: one method of feature extraciton from the numerical values.
Contact for help: gch_liu@163.com 
"""

import numpy as np
class mer_n:
    def __init__(self,ACGT,freq):
        self.ACGT=ACGT
        self.freq=freq
def kmer_finder(testseq,k):
    leng=len(testseq)

    kmer1=1
    if k>=kmer1:
        freq_A=(testseq.count('A'))/leng  #calculate the A frequency
        freq_C=(testseq.count('C'))/leng  #calculate the C frequency
        freq_G=(testseq.count('G'))/leng  #calculate the G frequency
        freq_T=(testseq.count('T'))/leng  #calculate the T frequency
        freq_GC=freq_G+freq_C  #calculate the GC content/frequency
        
        featuresAGCTGC=[freq_A,freq_C,freq_G,freq_T,freq_GC]

        features111 = np.array([],dtype=(str,10))
        features111 = np.resize(features111,[5,2])
        features111[:,0] = ['A','C','G','T','G_C']

        for i in range(0,5):
             features111[i,1] = featuresAGCTGC[i]
    else:
        return
# build a frequency product matrix, save frequency product like fG*fC
    f_ACGT=[freq_A,freq_C,freq_G,freq_T]
    kmer2=2
    if k>=kmer2:
        f_matrix2=np.ones((4,4))
        for i in range(0,4):
            for j in range(0,4):
                  f_matrix2[i,j]=f_ACGT[i]*f_ACGT[j]

        nucl='ACGT'
        
        freq112=np.ones((4,4))
        mer_2=[]
        for i in range(0,4):
           for j in range(0,4):
              freq112[i,j]=testseq.count(nucl[i]+nucl[j])/leng/f_matrix2[i,j]
              ACGT=nucl[i]+nucl[j]
              freq=freq112[i,j]
              mer_2.append(mer_n(ACGT,freq))

        # extract 
        nn22=len(mer_2)
#        ACGT_2_mer=np.ndarray((nn22,2),str)
        ACGT_2_mer=np.array([],dtype=(str,10))
        ACGT_2_mer = np.resize(ACGT_2_mer, [nn22, 2])
        freq_2_mer=np.zeros(nn22)
        for i in range(0,nn22):
           ACGT_2_mer[i,0]=mer_2[i].ACGT
           freq_2_mer[i]=mer_2[i].freq
    
        # combine
        for i in range(0,nn22):
           ACGT_2_mer[i,1] = freq_2_mer[i]

        features112=ACGT_2_mer
    else:
        features112=[]
#    %% 1.1.3 k=3, AAA/AAG/.../TTT frequency
#% build a frequency product matrix, save frequency product like fG*fC*fT
    kmer3=3
    if k>=kmer3:
        f_matrix3=np.ones((4,4,4))
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    f_matrix3[i,j,m]=f_ACGT[i]*f_ACGT[j]*f_ACGT[m]

        nucl='ACGT'
        freq113=np.ones((4,4,4))
        mer_3=[]
        for i in range(0,4):
           for j in range(0,4):
               for m in range(0,4):
                  freq113[i,j,m]=testseq.count(nucl[i]+nucl[j]+nucl[m])/leng/f_matrix3[i,j,m]
                  ACGT=nucl[i]+nucl[j]+nucl[m]
                  freq=freq113[i,j,m]
                  mer_3.append(mer_n(ACGT,freq))

        
        nn32=len(mer_3)
        ACGT_3_mer=np.array([],dtype=(str,10))
        ACGT_3_mer = np.resize(ACGT_3_mer,[nn32,2])
        freq_3_mer=np.zeros(nn32)
        for i in range(0,nn32):
           ACGT_3_mer[i,0]=mer_3[i].ACGT
           freq_3_mer[i]=mer_3[i].freq
        
        for i in range(0,nn32):
           ACGT_3_mer[i,1] = freq_3_mer[i]
        features113=ACGT_3_mer
    else:
        features113=[]


#% build a frequency product matrix, save frequency product like fG*fC*fT
    kmer4=4
    if k>=kmer4:
        f_matrix4=np.ones((4,4,4,4))
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    for n in range(0,4):
                        f_matrix4[i,j,m,n]=f_ACGT[i]*f_ACGT[j]*f_ACGT[m]*f_ACGT[n]

        freq114=np.ones((4,4,4,4))
        mer_4=[]
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    for n in range(0,4):
                        freq114[i,j,m,n]=testseq.count(nucl[i]+nucl[j]+nucl[m]+nucl[n])/leng/f_matrix4[i,j,m,n]
                        ACGT=nucl[i]+nucl[j]+nucl[m]+nucl[n]
                        freq=freq114[i,j,m,n]
                        mer_4.append(mer_n(ACGT,freq))

        #extract 
        nn42=len(mer_4)
        ACGT_4_mer=np.array([],dtype=(str,10))
        ACGT_4_mer = np.resize(ACGT_4_mer, [nn42, 2])
        freq_4_mer=np.zeros(nn42)
        for i in range(0,nn42):
           ACGT_4_mer[i,0]=mer_4[i].ACGT
           freq_4_mer[i]=mer_4[i].freq

        
#        % combin data and name
        for i in range(0,nn42):
           ACGT_4_mer[i,1] = freq_4_mer[i]

        features114=ACGT_4_mer
    else:
        features114=[]

#    % build a frequency product matrix, save frequency product like fG*fC*fT
    kmer5=5
    if k>=kmer5:
        f_matrix5=np.ones((4,4,4,4,4))
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    for n in range(0,4):
                        for t in range(0,4):
                            f_matrix5[i,j,m,n,t]=f_ACGT[i]*f_ACGT[j]*f_ACGT[m]*f_ACGT[n]*f_ACGT[t]

        freq115=np.ones((4,4,4,4,4))
        mer_5=[]
        for i in range(0,4):
           for j in range(0,4):
               for m in range(0,4):
                   for n in range(0,4):
                       for t in range(0,4):
                          freq115[i,j,m,n,t]=testseq.count(nucl[i]+nucl[j]+nucl[m]+nucl[n]+nucl[t])/leng/f_matrix5[i,j,m,n,t]
                          ACGT=nucl[i]+nucl[j]+nucl[m]+nucl[n]+nucl[t]
                          freq=freq115[i,j,m,n,t]
                          mer_5.append(mer_n(ACGT,freq))

        
#        %extract 
        nn52=len(mer_5)
        ACGT_5_mer=np.array([],dtype=(str,10))
        ACGT_5_mer = np.resize(ACGT_5_mer, [nn52, 2])
        freq_5_mer=np.zeros(nn52)
        for i in range(0,nn52):
           ACGT_5_mer[i,0]=mer_5[i].ACGT
           freq_5_mer[i]=mer_5[i].freq

        
#        % combin data and name
        for i in range(0,nn52):
           ACGT_5_mer[i,1] = freq_5_mer[i]

        features115=ACGT_5_mer
    else:
       features115=[]

       
#       % build a frequency product matrix, save frequency product like fG*fC*fT
    kmer6=6
    if k>=kmer6:
        f_matrix6=np.ones((4,4,4,4,4,4))
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    for n in range(0,4):
                        for t in range(0,4):
                            for s in range(0,4):
                                f_matrix6[i,j,m,n,t,s]=f_ACGT[i]*f_ACGT[j]*f_ACGT[m]*f_ACGT[n]*f_ACGT[t]*f_ACGT[s]

        freq116=np.ones((4,4,4,4,4,4))
        mer_6=[]
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    for n in range(0,4):
                        for t in range(0,4):
                            for s in range(0,4):
                                freq116[i,j,m,n,t,s]=testseq.count(nucl[i]+nucl[j]+nucl[m]+nucl[n]+nucl[t]+nucl[s])/leng/f_matrix6[i,j,m,n,t,s]
                                ACGT=nucl[i]+nucl[j]+nucl[m]+nucl[n]+nucl[t]+nucl[s]
                                freq=freq116[i,j,m,n,t,s]
                                mer_6.append(mer_n(ACGT,freq))


#        %extract 
        nn62=len(mer_6)
        ACGT_6_mer=np.array([],dtype=(str,10))
        ACGT_6_mer = np.resize(ACGT_6_mer,[nn62,2])
        freq_6_mer=np.zeros(nn62)
        for i in range(0,nn62):
           ACGT_6_mer[i,0]=mer_6[i].ACGT
           freq_6_mer[i]=mer_6[i].freq

#        % combin data and name
        for i in range(0,nn62):
           ACGT_6_mer[i,1] = freq_6_mer[i]

        features116=ACGT_6_mer
    else:
        features116=[]

    
#    % build a frequency product matrix, save frequency product like fG*fC*fT
    kmer7=7
    if k>=kmer7:
        f_matrix7=np.ones((4,4,4,4,4,4,4))
        for i in range(0,4):
            for j in range(0,4):
                for m in range(0,4):
                    for n in range(0,4):
                        for t in range(0,4):
                            for s in range(0,4):
                                for p in range(0,4):
                                    f_matrix7[i,j,m,n,t,s,p]=f_ACGT[i]*f_ACGT[j]*f_ACGT[m]*f_ACGT[n]*f_ACGT[t]*f_ACGT[s]*f_ACGT[p]

        freq117=np.ones((4,4,4,4,4,4,4))
        mer_7=[]
        
        for i in range(0,4):
           for j in range(0,4):
               for m in range(0,4):
                   for n in range(0,4):
                       for t in range(0,4):
                           for s in range(0,4):
                               for p in range(0,4):
                                  freq117[i,j,m,n,t,s,p]=testseq.count(nucl[i]+nucl[j]+nucl[m]+nucl[n]+nucl[t]+nucl[s]+nucl[p])/leng/f_matrix7[i,j,m,n,t,s,p]
                                  ACGT=nucl[i]+nucl[j]+nucl[m]+nucl[n]+nucl[t]+nucl[s]+nucl[p]
                                  freq=freq117[i,j,m,n,t,s,p]
                                  mer_7.append(mer_n(ACGT,freq))

#       extract 
        nn72=len(mer_7)
#        ACGT_7_mer=np.ndarray((nn72,2),str)
        ACGT_7_mer=np.array([],dtype=(str,10))
        ACGT_7_mer = np.resize(ACGT_7_mer, [nn72, 2])
        freq_7_mer=np.zeros(nn72)
        for i in range(0,nn72):
           ACGT_7_mer[i,0]=mer_7[i].ACGT
           freq_7_mer[i]=mer_7[i].freq

#        % combin data and name
        for i in range(0,nn72):
           ACGT_7_mer[i,1] = freq_7_mer[i]

        features117=ACGT_7_mer
    else:
        features117=[]

    if(k==3):
        kmer_ACGT_freq=np.vstack([features111,features112,features113])

    elif(k==4):
        kmer_ACGT_freq=np.vstack([features111,features112,features113,features114])

    elif(k==5):
        kmer_ACGT_freq=np.vstack([features111,features112,features113,features114,features115])

    elif(k==6):
        kmer_ACGT_freq=np.vstack([features111,features112,features113,features114,features115,features116])

    elif(k==7):
        kmer_ACGT_freq=np.vstack([features111,features112,features113,features114,features115,features116,features117])

    else:
        pass

    return [kmer_ACGT_freq,kmer_ACGT_freq[:,1]]

if __name__=='__main__':
    print(kmer_finder('ACGTCGT',3))