from random import randint, random
import matplotlib.pyplot as plt

def creation_echantillon(NA,NB,NC,ND) : # permet de générer l'échantillon ou sac
    echantillon = []
    for i in range(NA) :
        echantillon.append("A")
    for i in range(NA,NA+NB) :
        echantillon.append("B")
    for i in range(NA+NB,NA+NB+NC) :
        echantillon.append("C")
    for i in range(NA+NB+NC,NA+NB+NC+ND) :
        echantillon.append("D")
    return echantillon

def comptage_especes(sac,NbA, NbB, NbC, NbD): # fonction qui compte le nombre d'espèces
    l = len(sac)
    NbeA, NbeB, NbeC, NbeD = 0,0,0,0
    for i in range(l):
        if sac[i]=="A" :
            NbeA = NbeA + 1
        elif sac[i]=="B" :
            NbeB = NbeB + 1
        elif sac[i]=="C" :
            NbeC = NbeC + 1
        elif sac[i]=="D" :
            NbeD = NbeD + 1
    NbA.append(NbeA)
    NbB.append(NbeB)
    NbC.append(NbeC)
    NbD.append(NbeD)
    return NbA, NbB, NbC, NbD

def collision(sac, collisionAB) : # fonction qui génère des collisions et remplace A, B éventuellement
    l = len(sac)
    N1 = randint(0,l-1)
    N2 = randint(0,l-1)
    while N2 == N1 :
        N2 = randint(0,l-1)
    if (sac[N1] == "A" and sac[N2] == "B") or (sac[N1] == "B" and sac[N2] == "A") :
        x=random()
        if x > collisionAB :
            sac[N1], sac[N2] = "C", "D"
    return sac

# programme principal
NA = 1000  # nombre d'espèces A pouvant être modifié
NB = 800    # nombre d'espèces B pouvant être modifié
NC = 0  # nombre d'espèces C pouvant être modifié
ND = 100  # nombre d'espèces D pouvant être modifié
collisionAB = 0.05  # facteur de collision entre A et B compris entre 0 et 1 pouvant être modifié, plus il est grand moins les chocs sont efficaces
# pas de collision entre C et D comme demandé par Michel
temps = 10000  # variable symbolisant le temps
NbA, NbB, NbC, NbD = [], [], [], [] # création de liste du nombre d'espèces A, B, C et D
sac = creation_echantillon(NA,NB,NC,ND)  # fonction qui crée l'échantillon
t=[0] # liste de temps qui a initialement la valeur 0
comptage_especes(sac, NbA, NbB, NbC, NbD) # comptage des espèces à t = 0
for i in range(temps) : # on fait une boucle pour 1000 unités de temps
    collision(sac, collisionAB) # on envisage une collision à chaque unité de temps
    comptage_especes(sac, NbA, NbB, NbC, NbD) # à chaque collision, on compte les espèces
    t.append(i+1) # on stocke dans une liste les différents temps
# affichage des courbes
plt.grid()
plt.title("Population des espèces A, B, C et D")
plt.xlabel("Nombre de Chocs")
plt.ylabel("Nombre d'entités")
plt.plot(t,NbA,"r--+",label="Nombre d'espèces A") # on trace la courbe Nbre de A en fonction de t
plt.plot(t,NbB,"b--*",label="Nombre d'espèces B") # idem pour B
plt.plot(t,NbC,"g--o",label="Nombre d'espèces C") # idem pour C
plt.plot(t,NbD,"k--+",label="Nombre d'espèces D") # idemen pour D
plt.legend()
plt.show()