import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.backends.backend_pdf import PdfPages
import numpy.random as rd
import matplotlib
from matplotlib.animation import PillowWriter, FFMpegWriter
matplotlib.use('Qt5Agg')

# Paramètres de la simulation
N = 1000  # Nombre de particules
T = 1000  # Nombre de pas de calcul
l = 200  # Taille de la simulation
dt = 0.05  # Pas de temps simulation
a = 1  # Pas
pos = np.array([[0, 0]] * N)  # Tableau des positions de chaque point
step = np.linspace(0, T, T)  # Tableau de T valeurs
mx = np.zeros(T)  # Le tableau qui contient la moyenne des positions
my = np.zeros(T)
stdx = np.zeros(len(step))  # Le tableau qui contient l'écart-type des positions
stdy = np.zeros(len(step))

# Création du dossier pour les images
if not os.path.exists(f"images_N{N}"):
    os.makedirs(f"images_N{N}")

def actualise_position(M):
    p = [(M[0], M[1] + a), (M[0], M[1] - a), (M[0] + a, M[1]), (M[0] - a, M[1])]  
    return p[rd.choice([0, 1, 2, 3])]

# Paramètres graphiques
plt.ion()  # Mode interactif activé pour Spyder
fig = plt.figure(figsize=(14, 24))  # Taille de la figure

# Ajout des sous-graphes (3 lignes, 2 colonnes) avec des espacements ajustés
ax1 = fig.add_subplot(3, 2, (1, 2))  # Le premier graphique occupe toute la largeur
ax2 = fig.add_subplot(3, 2, 3)  # Deuxième ligne, première colonne
ax3 = fig.add_subplot(3, 2, 4)  # Deuxième ligne, deuxième colonne
ax4 = fig.add_subplot(3, 2, 5)  # Troisième ligne, première colonne
ax5 = fig.add_subplot(3, 2, 6)  # Troisième ligne, deuxième colonne

# Ajout du titre principal
fig.suptitle(f"Simulation de diffusion pour N = {N} particules", fontsize=20, y=0.97)

# Réorganisation des titres
ax1.set_title("Position des particules", fontsize=16)
ax2.set_title("Moyenne des positions $<x>$", fontsize=14)
ax3.set_title("Écart-type des positions $<x^2>$", fontsize=14)
ax4.set_title("Moyenne des positions $<y>$", fontsize=14)
ax5.set_title("Écart-type des positions $<y^2>$", fontsize=14)

# Configuration des axes
ax1.set_xlim(-l // 2, l // 2)
ax1.set_ylim(-l // 2, l // 2)
ax1.set_xlabel("x", fontsize=12)
ax1.set_ylabel("y", fontsize=12)
ax2.set_ylabel("$<x>$", fontsize=12)
ax2.set_xlabel("$t$", fontsize=12)
ax3.set_ylabel("$<x^2>$", fontsize=12)
ax3.set_xlabel("$t$", fontsize=12)
ax4.set_ylabel("$<y>$", fontsize=12)
ax4.set_xlabel("$t$", fontsize=12)
ax5.set_ylabel("$<y^2>$", fontsize=12)
ax5.set_xlabel("$t$", fontsize=12)

# Assurer un aspect carré pour le graphique 1 (Position des particules)
ax1.set_aspect('equal', adjustable='box')

# Grilles
for axis in [ax1, ax2, ax3, ax4, ax5]:
    axis.grid(True)

# Ajustement des espacements entre les sous-graphiques
fig.subplots_adjust(hspace=0.4, wspace=0.2)

# Initialisation des graphiques
points, = ax1.plot([], [], 'bo', ms=3)
linex, = ax2.plot([], [], c='r')
line_stdx, = ax3.plot([], [], c='g')
liney, = ax4.plot([], [], c='r')
line_stdy, = ax5.plot([], [], c='g')

def init():
    points.set_data([], [])
    linex.set_data([], [])
    line_stdx.set_data([], [])
    liney.set_data([], [])
    line_stdy.set_data([], [])
    return points, linex, line_stdx, liney, line_stdy

def animate(i):
    global pos
    pos = np.array([actualise_position(M) for M in pos])
    points.set_data(pos[:, 0], pos[:, 1])
    mx[i] = np.mean(pos[:, 0])
    my[i] = np.mean(pos[:, 1])
    stdx[i] = np.var(pos[:, 0])
    stdy[i] = np.var(pos[:, 1])
    linex.set_data(step[:i] * dt, mx[:i])
    line_stdx.set_data(step[:i] * dt, stdx[:i])
    liney.set_data(step[:i] * dt, my[:i])
    line_stdy.set_data(step[:i] * dt, stdy[:i])

    # Ajuster les limites pour chaque graphique
    ax2.set_xlim(0, (i + 1) * dt)
    ax2.set_ylim(min(mx[:i + 1]) - 0.5, max(mx[:i + 1]) + 0.5)
    ax3.set_xlim(0, (i + 1) * dt)
    ax3.set_ylim(0, max(stdx[:i + 1]) + 0.5)
    ax4.set_xlim(0, (i + 1) * dt)
    ax4.set_ylim(min(my[:i + 1]) - 0.5, max(my[:i + 1]) + 0.5)
    ax5.set_xlim(0, (i + 1) * dt)
    ax5.set_ylim(0, max(stdy[:i + 1]) + 0.5)

    # Sauvegarder une image toutes les 10 itérations
    if i % 10 == 0:
        fig.savefig(f"images_N{N}/step_N{N}_{i:04d}.png", bbox_inches='tight')

    return points, linex, line_stdx, liney, line_stdy

# Modifier la taille pour être divisible par 2
width, height = fig.get_size_inches()
width, height = int(width * fig.dpi), int(height * fig.dpi)
if height % 2 != 0:
    height += 1
if width % 2 != 0:
    width += 1

fig.set_size_inches(width / fig.dpi, height / fig.dpi)

# Créer l'animation
ani = FuncAnimation(fig, animate, frames=T, init_func=init, interval=dt * 1000, repeat=False, blit=False)

# Enregistrer d'abord l'animation dans le format GIF
gif_writer = PillowWriter(fps=30)
ani.save(f"diffusion_N{N}.gif", writer=gif_writer)

# Ensuite, enregistrer dans le format MP4
mp4_writer = FFMpegWriter(fps=30)
ani.save(f"diffusion_N{N}.mp4", writer=mp4_writer)

plt.show(block=True)

# Sauvegarde des images dans un PDF
with PdfPages(f"animation_images_N{N}.pdf") as pdf:
    for image_file in sorted(os.listdir(f"images_N{N}")):
        if image_file.endswith(".png"):
            img = plt.imread(f"images_N{N}/{image_file}")
            fig_img, ax = plt.subplots(figsize=(12, 8))
            ax.imshow(img)
            ax.axis('off')
            pdf.savefig(fig_img, dpi=300, bbox_inches='tight')
            plt.close(fig_img)

import os
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import numpy.random as rd
import matplotlib
matplotlib.use('Qt5Agg')

# Paramètres de la simulation
N = 1000  # Nombre de particules
T = 1000  # Nombre de pas de calcul
l = 200  # Taille de la simulation
dt = 0.05  # Pas de temps simulation
a = 1  # Pas
pos = np.array([[0, 0]] * N)  # Tableau des positions de chaque point
step = np.linspace(0, T, T)  # Tableau de T valeurs
mx = np.zeros(T)  # Le tableau qui contient la moyenne des positions
my = np.zeros(T)
stdx = np.zeros(len(step))  # Le tableau qui contient l'écart-type des positions
stdy = np.zeros(len(step))

# Création du dossier pour les images
if not os.path.exists(f"images_N{N}"):
    os.makedirs(f"images_N{N}")

def actualise_position(M):
    p = [(M[0], M[1] + a), (M[0], M[1] - a), (M[0] + a, M[1]), (M[0] - a, M[1])]  
    return p[rd.choice([0, 1, 2, 3])]

# Paramètres graphiques
plt.ion()  # Mode interactif activé pour Spyder
fig = plt.figure(figsize=(14, 24))  # Taille de la figure

# Ajout des sous-graphes (3 lignes, 2 colonnes) avec des espacements ajustés
ax1 = fig.add_subplot(3, 2, (1, 2))  # Le premier graphique occupe toute la largeur
ax2 = fig.add_subplot(3, 2, 3)  # Deuxième ligne, première colonne
ax3 = fig.add_subplot(3, 2, 4)  # Deuxième ligne, deuxième colonne
ax4 = fig.add_subplot(3, 2, 5)  # Troisième ligne, première colonne
ax5 = fig.add_subplot(3, 2, 6)  # Troisième ligne, deuxième colonne

# Ajout du titre principal
fig.suptitle(f"Simulation de diffusion pour N = {N} particules", fontsize=20, y=0.97)

# Réorganisation des titres
ax1.set_title("Position des particules", fontsize=16)
ax2.set_title("Moyenne des positions $<x>$", fontsize=14)
ax3.set_title("Écart-type des positions $<x^2>$", fontsize=14)
ax4.set_title("Moyenne des positions $<y>$", fontsize=14)
ax5.set_title("Écart-type des positions $<y^2>$", fontsize=14)

# Configuration des axes
ax1.set_xlim(-l // 2, l // 2)
ax1.set_ylim(-l // 2, l // 2)
ax1.set_xlabel("x", fontsize=12)
ax1.set_ylabel("y", fontsize=12)
ax2.set_ylabel("$<x>$", fontsize=12)
ax2.set_xlabel("$t$", fontsize=12)
ax3.set_ylabel("$<x^2>$", fontsize=12)
ax3.set_xlabel("$t$", fontsize=12)
ax4.set_ylabel("$<y>$", fontsize=12)
ax4.set_xlabel("$t$", fontsize=12)
ax5.set_ylabel("$<y^2>$", fontsize=12)
ax5.set_xlabel("$t$", fontsize=12)

# Assurer un aspect carré pour le graphique 1 (Position des particules)
ax1.set_aspect('equal', adjustable='box')

# Grilles
for axis in [ax1, ax2, ax3, ax4, ax5]:
    axis.grid(True)

# Ajustement des espacements entre les sous-graphiques
fig.subplots_adjust(hspace=0.4, wspace=0.2)

# Initialisation des graphiques
points, = ax1.plot([], [], 'bo', ms=3)
linex, = ax2.plot([], [], c='r')
line_stdx, = ax3.plot([], [], c='g')
liney, = ax4.plot([], [], c='r')
line_stdy, = ax5.plot([], [], c='g')

def init():
    points.set_data([], [])
    linex.set_data([], [])
    line_stdx.set_data([], [])
    liney.set_data([], [])
    line_stdy.set_data([], [])
    return points, linex, line_stdx, liney, line_stdy

def animate(i):
    global pos
    pos = np.array([actualise_position(M) for M in pos])
    points.set_data(pos[:, 0], pos[:, 1])
    mx[i] = np.mean(pos[:, 0])
    my[i] = np.mean(pos[:, 1])
    stdx[i] = np.var(pos[:, 0])
    stdy[i] = np.var(pos[:, 1])
    linex.set_data(step[:i] * dt, mx[:i])
    line_stdx.set_data(step[:i] * dt, stdx[:i])
    liney.set_data(step[:i] * dt, my[:i])
    line_stdy.set_data(step[:i] * dt, stdy[:i])

    ax2.set_xlim(0, (i + 1) * dt)
    ax2.set_ylim(min(mx[:i + 1]) - 0.5, max(mx[:i + 1]) + 0.5)
    ax3.set_xlim(0, (i + 1) * dt)
    ax3.set_ylim(0, max(stdx[:i + 1]) + 0.5)
    ax4.set_xlim(0, (i + 1) * dt)
    ax4.set_ylim(min(my[:i + 1]) - 0.5, max(my[:i + 1]) + 0.5)
    ax5.set_xlim(0, (i + 1) * dt)
    ax5.set_ylim(0, max(stdy[:i + 1]) + 0.5)

    return points, linex, line_stdx, liney, line_stdy

ani = FuncAnimation(fig, animate, frames=T, init_func=init, interval=dt * 1000, repeat=False, blit=False)
plt.show(block=True)