import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# --- Demande à l'utilisateur ---
user_input = input("GIF ? (o/n) : ").strip().lower()
GENERATE_GIF = user_input == 'o'

# --- Fichiers ---
ref_file = "constellation_ref.dat"   # constellation de référence
rx_file = "constellation.dat"        # symboles corrigés par la PLL

# Charger les données
ref_constel = np.loadtxt(ref_file)  # colonnes I Q
rx_data = np.loadtxt(rx_file)       # colonnes I Q

# Paramètres
window_size = 50   # nombre de symboles à afficher dans la fenêtre
speed_factor = 5   # pour accélérer l'animation
N_frames = len(rx_data)

# Figure
fig, ax = plt.subplots(figsize=(6,6))
ax.set_title("Rolling Constellation avec PLL")
ax.set_xlabel("I")
ax.set_ylabel("Q")
ax.grid(True)

# Constellation de référence
ax.scatter(ref_constel[:,0], ref_constel[:,1], c='red', marker='x', label='Référence')

# Points PLL
scat = ax.scatter([], [], c='blue', s=20, label='Points récents')
last_point = ax.scatter([], [], c='green', s=50, label='Dernier symbole')

ax.set_xlim(np.min(ref_constel[:,0])-0.5, np.max(ref_constel[:,0])+0.5)
ax.set_ylim(np.min(ref_constel[:,1])-0.5, np.max(ref_constel[:,1])+0.5)
ax.legend()

# Initialisation
def init():
    scat.set_offsets(np.empty((0,2)))
    last_point.set_offsets(np.empty((0,2)))
    return scat, last_point

# Mise à jour
def update(frame):
    idx = min(frame*speed_factor, N_frames)
    start_idx = max(0, idx - window_size)
    data_window = rx_data[start_idx:idx]
    
    scat.set_offsets(data_window)
    if len(data_window) > 0:
        last_point.set_offsets(data_window[-1:])
    return scat, last_point

# Création de l'animation
ani = FuncAnimation(fig, update, frames=int(N_frames/speed_factor)+1,
                    init_func=init, blit=True, interval=20)

# --- Sauvegarde GIF conditionnelle ---
if GENERATE_GIF:
    ani.save("pll_constellation.gif", writer='pillow', fps=30)
    print("GIF généré : pll_constellation.gif")

# Affichage à l'écran
plt.show()