import numpy as np
import sys
from typing import Tuple

if len(sys.argv) >= 2:
    u0 = int(sys.argv[1])
else:
    u0 = 115

print("u0 = ", u0)
print()

## Partie 1 - Génération de séquences ADN aléatoires

## Question 1
u = np.zeros(100000, dtype=np.int64)
u[0] = u0
for k in range(1, 100000):
    u[k] = (15091 * u[k-1]) % 64007

print("Question 1")
print("u10 = ", u[10])
print("u100 = ", u[100])
print("u1000 = ", u[1000])
print()

## Question 2
def S(k : int, l : int) -> str:
    res = ""
    for i in range(l):
        x = u[i + k]
        if x % 4 == 0:
            lettre = "A"
        elif x % 4 == 1:
            lettre = "T"
        elif x % 4 == 2:
            lettre = "G"
        else:
            lettre = "C"
        res = res + lettre
    return res

print("Question 2")
print("a/ S(0, 5) = ", S(0, 5))
print("b/ S(10, 5) = ", S(10, 5))
print("c/ S(20, 5) = ", S(20, 5))
print()

## Partie 2 - Recherche de motif dans une séquence

## Question 3 : recherche exacte

def recherche_exacte(seq : str, motif : str) -> list:
    positions = [] # positions où le motif a été détecté
    n = len(seq)
    m = len(motif)
    for i in range(n - m + 1): #i est la position de la fenetre
        k = 0 #k est la positition dans le motif
        while (k < m and seq[i + k] == motif[k]):
            k = k + 1
        if k == m: # si on est arrive au bout du motif
            positions.append(i) # alors il y a occurrence
    return positions

print("Question 3")
print("a/ ", recherche_exacte("ATATATACATATA", "TATA"))
print("b/ ", recherche_exacte(S(20, 1000), "TATA"))
l = recherche_exacte(S(50, 10000), "TATA")
print("c/ ", len(l))
print()

## Question 4 : recherche avec substitutions

#     A  G  C  T
# A   0  2  6  8
# G   2  0  9  6
# C   6  9  0  1
# T   8  6  1  0

matsub = {}
matsub['A', 'A'] = 0
matsub['A', 'G'] = 2
matsub['A', 'C'] = 6
matsub['A', 'T'] = 8

matsub['G', 'A'] = 2
matsub['G', 'G'] = 0
matsub['G', 'C'] = 9
matsub['G', 'T'] = 6

matsub['C', 'A'] = 6
matsub['C', 'G'] = 9
matsub['C', 'C'] = 0
matsub['C', 'T'] = 1

matsub['T', 'A'] = 8
matsub['T', 'G'] = 6
matsub['T', 'C'] = 1
matsub['T', 'T'] = 0

def recherche_subs(seq : str, motif : str, coutmax : int) -> list:
    positions = [] # positions où le motif a été détecté
    n = len(seq)
    m = len(motif)
    for i in range(n - m + 1): #i est la position de la fenetre
        cout = 0 #represente le cout de l'occurrence
        for k in range(m):
            cout += matsub[seq[i + k], motif[k]]
        if cout <= coutmax:
            positions.append(i)
    return positions


print("Question 4")
print("a/ ", recherche_subs("ATATATACATATA", "TATA", 3))
print("b/ ", (recherche_subs(S(50, 10000), "TATA", 3))[:5])
s = S(50, 10000)
l = recherche_subs(s, "GATTACA", 3)
print("c/  ", len(l))

## Partie 3 - Alignement de séquences ADN

# Le cout d'une substitution est donnée par la matrice matsub
# Le cout d'une suppression est 4
# Le cout d'une insertion est 5

cout_supp = 5 
cout_insert = 4 

## Question 5 : score d'un alignement

def score_alignement(s1 : str, s2 : str) -> int:
    assert len(s1) == len(s2) # Précondition
    n = len(s1)
    score = 0
    for i in range(n):
        if not(s1[i] in ['A', 'G', 'C', 'T', '-']):
            raise ValueError("Sequence invalide")
        if not(s2[i] in ['A', 'G', 'C', 'T', '-']):
            raise ValueError("Sequence invalide")
        if (s1[i] == '-' and s2[i] == '-'):
            raise ValueError("Alignement inccorect")
        if (s1[i] == '-'):
            score += cout_insert
        elif (s2[i] == '-'):
            score += cout_supp
        else:
            score += matsub[s1[i], s2[i]]
    return score

print("Question 5")
s1 = "AATTGCAT"
s2 = "AA--GCAT"
print("a/ ", score_alignement(s1, s2))
s1 = "A-TTGCA-"
s2 = "AAT--CAT"
print("b/ ", score_alignement(s1, s2))
s1 = S(2000, 50)
s2 = S(3000, 50)
print("c/ ", score_alignement(s1, s2))
print()

print("Question 6 : réponse C")
print()

## Question 7 et 8 : score d'un alignement optimal

# On notera S(p, q) le score d'un alignement optimal entre le préfixe de longueur p de s1
# et le préfixe de longueur q de s2

# Dans la question 7 l'étudiant trouve la relation de récurrence adaptée eu probleme
# Dans la question 8 l'étudiant calcule par programmation dynamique le score optimal sans
# construire l'alignement correspondant

def score_optimal(s1 : str, s2 : str) -> int:
    n = len(s1)
    m = len(s2)
    T = np.zeros([n+1,m+1], dtype='int')
    
    # Initialisation
    for j in range(m+1):
        T[0, j] = j * cout_insert
    for i in range(n+1):
        T[i, 0] = i * cout_supp
    
    # Remlissage
    for i in range(1, n+1):
        for j in range(1, m+1):
            # Cout dans le cas de la substitution des derniers nucleotides
            subs = matsub[s1[i-1], s2[j-1]] + T[i-1, j-1]
            # Cout dans le cas de la suppression du dernier nucleotide de s1
            supp = cout_supp + T[i-1, j]
            # Cout dans le cas de l'insertion du dernier nucleotide de s2
            insert = cout_insert + T[i, j-1]
            # On garde la solution qui produit le meilleure score
            T[i, j] = min(subs, supp, insert)
    return T[n, m]

print("Question 7")
s1 = "GGTTCA"
s2 = "GGTCA"
print("a/ ", score_optimal(s1, s2))
s1 = "CATTCACATCTTTAGCA"
s2 = "TTTCAGATCTCTATGCA"
print("b/ ", score_optimal(s1, s2))
s1 = S(2000,50)
s2 = S(3000,50)
print("c/ ", score_optimal(s1, s2))
print()

# Question 8

# La meme chose mais on reconstruit l'alignement
# Pour cela dans la programmation dynamique, on se sert d'une 3e matrice
# qui contient :
# 1 si une subsititon a ete faite
# 2 si une suppression a ete realisee
# 3 si une insertion a ete choisie


def alignement_optimal(s1 : str, s2 : str) -> Tuple[str, str]:
    n = len(s1)
    m = len(s2)
    T = np.zeros([n+1,m+1], dtype='int')
    U = np.zeros([n+1,m+1], dtype='int')
    # Initialisation
    for j in range(m+1):
        T[0, j] = j * cout_insert
    for i in range(n+1):
        T[i, 0] = i * cout_supp
    
    # Remlissage
    for i in range(1, n+1):
        for j in range(1, m+1):
            # Cout dans le cas de la substitution des derniers nucleotides
            subs = matsub[s1[i-1], s2[j-1]] + T[i-1, j-1]
            # Cout dans le cas de la suppression du dernier nucleotide de s1
            supp = cout_supp + T[i-1, j]
            # Cout dans le cas de l'insertion du dernier nucleotide de s2
            insert = cout_insert + T[i, j-1]
            # On garde la solution qui produit le meilleure score
            # en sauvegardant dans U l'operation choisie
            if subs <= supp and subs <= insert:
                T[i, j] = subs
                U[i, j] = 1
            elif supp <= subs and supp <= insert:
                T[i, j] = supp
                U[i, j] = 2
            else:
                T[i, j] = insert
                U[i, j] = 3

    # Reconstruction de l'alignement optimal
    # a partir des informations memorisees
    r1 = ""
    r2 = ""
    i, j = n, m
    while (i > 0 and j > 0):
        if U[i,j] == 1:
            r1 = s1[i - 1] + r1
            r2 = s2[j - 1] + r2
            i = i - 1
            j = j - 1
        elif U[i, j] == 2:
            r1 = s1[i - 1] + r1
            r2 = "-" + r2
            i = i - 1
        else:
            r1 = "-" + r1
            r2 = s2[j - 1] + r2
            j = j - 1
    if i == 0:
        r1 = j * "-" + r1
        r2 = s2[0:j] + r2
    else:
        r1 = s1[0:i] + r1
        r2 = i * "-" + r2
    
    return (r1, r2)
    
print("Question 8")
s1 = "GGTTCA"
s2 = "GGTCA"
(r1, r2) = alignement_optimal(s1, s2)
print("a/ ", r1, " et ", r2)
s1 = "CATTCACATCTTTAGCA"
s2 = "TTTCAGATCTCTATGCA"
(r1, r2) = alignement_optimal(s1, s2)
print("b/ ", r1, " et ", r2)
s1 = S(2000,50)
s2 = S(3000,50)
(r1, r2) = alignement_optimal(s1, s2)
print("c/ ", r1, " et ", r2)
print()