import numpy as np
import matplotlib.pyplot as plt

"""
La suite logistique

x_{n+1} = f(x_n) ou f(x) = cx(1 -x)
avec c dans l'intervalle [0, 4]

On cree une animation (rudimentaire) sur l'evolution de 
l'escargot defini par les N premiers termes de la suite 
logistique en fonction de la constante c.
"""

def suiteLogistique(c, x0, N):
    """
    Parameters
    ----------
    c : reel entre 0 et 4, ou vector
    x0 : terme initial
    N : entier > 0

    Returns
    -------
    si c est un reel, on pose d=1
    si c est un vecteur, on pose d=len(c)
    On retourne la matrice de taille (N+1, d) ayant sur chaque 
    colonne j les termes x_0, x_1, ..., x_N correspondant a la 
    constante c[j].
    """
    if type(c)==float or type(c)==int:
        d = 1
    else:
        d = len(c)
        
    x = x0*np.ones(d)
    res = np.zeros((N+1, d))
    res[0, :] = x
    for k in range(1, N+1):
        x = c*x*(1 - x)
        res[k, :] = x
        
    return res
        

def pourLEscargot_animation(c, x0, N):
    """
    Parameters
    ----------
    c : vector
    x0 : terme initial
    N : entier > 0

    Returns
    -------
    On retourne les abscisses et les ordonnees sous forme
    de matrices (2N, len(c)) ; sur chaque colonne j on a les 
    abscisses et ordonnees de la ligne  polygonale determinee 
    par 
    [x_0, 0], [x_0, x_1], [x_1, x_1], [x_1, x_2], ..., [x_{n-1}, x_n]
    correspondant a la constante c[j].
    """
    A = suiteLogistique(c, x0, N)
    X = np.zeros((2*N, len(c)))
    Y = np.zeros((2*N, len(c)))
    
    for k in range(N):
        X[2*k, :] = A[k, :]
        X[2*k + 1, :] = A[k, :]
        
    Y[1:-1, :] = X[2:, :]
    Y[-1, :] = A[-1, :]
    return X, Y


###################################################################
#### donnees initiales
c = np.linspace(1.5, 4, 120)
x0 = .14
N = 300

X, Y = pourLEscargot_animation(c, x0, N)
print(np.shape(X), np.shape(Y))

delta_t = .2  # proportionnel a l'intervalle de temps de vue
              # de chaque image
plt.ion()  # interactive is on
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1, aspect="equal")
ax.set_xlim(0, 1.)
ax.set_ylim(0, 1.)


### premiere image
k = 0
x = np.linspace(0, 1, 100)
y_parabole = c[k]*x*(1 - x)
c_3 = round(c[k], 3)

ax.set_title(f"les premiers {N} termes pour c = {c_3}")
ax.plot([0, 1], [0, 1], c="r", linewidth=.5)
parabole, = ax.plot(x, y_parabole, c="b", linewidth=.5)
escargot, = ax.plot(X[:, k], Y[:, k], c="k", linewidth=1.1)
plt.pause(delta_t)


### animation
for k in range(len(c)):
   y_parabole = c[k]*x*(1 - x)
   c_3 = round(c[k], 3)
   ax.set_title(f"les premiers {N} termes pour c = {c_3}")
   parabole.set_ydata(y_parabole)
   escargot.set_xdata(X[:, k])
   escargot.set_ydata(Y[:, k])
   plt.pause(delta_t)
   plt.show()