import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.widgets import Button

def animar_trayectorias_combinadas(x, y, z, t, wc=None, wz=None, skip_points=5):
    """
    Muestra animaciones 2D (XY) y 3D en una misma ventana (subplots horizontales).
    
    Parámetros:
    -----------
    x, y, z : array-like
        Arrays con las coordenadas de la trayectoria.
    t : array-like
        Array con los tiempos correspondientes.
    wc, wz : float, optional
        Parámetros para el título.
    skip_points : int, optional
        Saltar puntos para optimizar rendimiento (default=5).
    """
    # Configuración de la figura
    fig = plt.figure(figsize=(16, 6))
    
    # ---- Subplot 1: Animación 2D (XY) ----
    ax_2d = fig.add_subplot(121)
    ax_2d.set_xlim(min(x)-0.1, max(x)+0.1)
    ax_2d.set_ylim(min(y)-0.1, max(y)+0.1)
    ax_2d.set_xlabel('$x$', fontsize=12)
    ax_2d.set_ylabel('$y$', fontsize=12)
    title_2d = 'Trayectoria 2D (XY)'
    if wc is not None and wz is not None:
        title_2d += f'\n$\omega_c = {wc}$, $\omega_z = {wz}$'
    ax_2d.set_title(title_2d, fontsize=14)
    ax_2d.grid(True)

    line_2d, = ax_2d.plot([], [], 'b-', lw=1, alpha=0.6)
    point_2d, = ax_2d.plot([], [], 'ro', markersize=6)

    # ---- Subplot 2: Animación 3D ----
    ax_3d = fig.add_subplot(122, projection='3d')
    ax_3d.set_xlim(min(x), max(x))
    ax_3d.set_ylim(min(y), max(y))
    ax_3d.set_zlim(min(z), max(z))
    ax_3d.set_xlabel('$x$', fontsize=12)
    ax_3d.set_ylabel('$y$', fontsize=12)
    ax_3d.set_zlabel('$z$', fontsize=12)
    title_3d = 'Trayectoria 3D'
    if wc is not None and wz is not None:
        title_3d += f'\n$\omega_c = {wc}$, $\omega_z = {wz}$'
    ax_3d.set_title(title_3d, fontsize=14)

    line_3d, = ax_3d.plot([], [], [], 'b-', lw=1, alpha=0.6)
    point_3d, = ax_3d.plot([], [], [], 'ro', markersize=6)

    # Función de inicialización
    def init():
        line_2d.set_data([], [])
        point_2d.set_data([], [])
        
        line_3d.set_data([], [])
        line_3d.set_3d_properties([])
        point_3d.set_data([], [])
        point_3d.set_3d_properties([])
        
        return line_2d, point_2d, line_3d, point_3d

    # Función de actualización
    def update(frame):
        # Actualiza 2D
        line_2d.set_data(x[:frame*skip_points], y[:frame*skip_points])
        point_2d.set_data(x[frame*skip_points], y[frame*skip_points])
        
        # Actualiza 3D
        line_3d.set_data(x[:frame*skip_points], y[:frame*skip_points])
        line_3d.set_3d_properties(z[:frame*skip_points])
        point_3d.set_data(x[frame*skip_points], y[frame*skip_points])
        point_3d.set_3d_properties(z[frame*skip_points])
        ax_3d.view_init(elev=20, azim=frame/2)
        
        return line_2d, point_2d, line_3d, point_3d

    # Calcula frames
    frames = len(t) // skip_points
    
    # Crea la animación
    ani = FuncAnimation(fig, update, frames=frames,
                       init_func=init, blit=True, interval=30)

    plt.tight_layout()
    return ani

def animar_trayectorias_con_boton(x, y, z, t, wc=None, wz=None, skip_points=5):
    """
    Muestra animaciones 2D (XY) y 3D en una misma ventana con botón Play/Pause.
    """
    # Configuración de la figura
    fig = plt.figure(figsize=(16, 7))
    
    # ---- Subplot 1: Animación 2D (XY) ----
    ax_2d = fig.add_subplot(121)
    ax_2d.set_xlim(min(x)-0.1, max(x)+0.1)
    ax_2d.set_ylim(min(y)-0.1, max(y)+0.1)
    ax_2d.set_xlabel('$x$', fontsize=12)
    ax_2d.set_ylabel('$y$', fontsize=12)
    title_2d = 'Trayectoria 2D (XY)'
    if wc is not None and wz is not None:
        title_2d += f'\n$\omega_c = {wc}$, $\omega_z = {wz}$'
    ax_2d.set_title(title_2d, fontsize=14)
    ax_2d.grid(True)

    line_2d, = ax_2d.plot([], [], 'b-', lw=1, alpha=0.6)
    point_2d, = ax_2d.plot([], [], 'ro', markersize=6)

    # ---- Subplot 2: Animación 3D ----
    ax_3d = fig.add_subplot(122, projection='3d')
    ax_3d.set_xlim(min(x), max(x))
    ax_3d.set_ylim(min(y), max(y))
    ax_3d.set_zlim(min(z), max(z))
    ax_3d.set_xlabel('$x$', fontsize=12)
    ax_3d.set_ylabel('$y$', fontsize=12)
    ax_3d.set_zlabel('$z$', fontsize=12)
    title_3d = 'Trayectoria 3D'
    if wc is not None and wz is not None:
        title_3d += f'\n$\omega_c = {wc}$, $\omega_z = {wz}$'
    ax_3d.set_title(title_3d, fontsize=14)

    line_3d, = ax_3d.plot([], [], [], 'b-', lw=1, alpha=0.6)
    point_3d, = ax_3d.plot([], [], [], 'ro', markersize=6)

    # ---- Botón Play/Pause ----
    ax_button = plt.axes([0.4, 0.02, 0.2, 0.05])  # Posición del botón [x, y, ancho, alto]
    button = Button(ax_button, 'Pause', color='lightgoldenrodyellow', hovercolor='0.8')

    # Función de inicialización
    def init():
        line_2d.set_data([], [])
        point_2d.set_data([], [])
        line_3d.set_data([], [])
        line_3d.set_3d_properties([])
        point_3d.set_data([], [])
        point_3d.set_3d_properties([])
        return line_2d, point_2d, line_3d, point_3d

    # Función de actualización
    def update(frame):
        line_2d.set_data(x[:frame*skip_points], y[:frame*skip_points])
        point_2d.set_data(x[frame*skip_points], y[frame*skip_points])
        line_3d.set_data(x[:frame*skip_points], y[:frame*skip_points])
        line_3d.set_3d_properties(z[:frame*skip_points])
        point_3d.set_data(x[frame*skip_points], y[frame*skip_points])
        point_3d.set_3d_properties(z[frame*skip_points])
        ax_3d.view_init(elev=20, azim=frame/2)
        return line_2d, point_2d, line_3d, point_3d

    # Crea la animación
    frames = len(t) // skip_points
    ani = FuncAnimation(fig, update, frames=frames, init_func=init, 
                        blit=True, interval=30, repeat=True)

    # Función para pausar/reanudar
    def toggle_animation(event):
        if ani.event_source:
            ani.event_source.stop()
            button.label.set_text('Play')
        else:
            ani.event_source.start()
            button.label.set_text('Pause')
        fig.canvas.draw()

    button.on_clicked(toggle_animation)

    plt.tight_layout()
    plt.show()
    return ani

def animar_trayectorias_inicio_pausa(x, y, z, t, wc=None, wz=None, skip_points=5):
    """
    Muestra animaciones 2D (XY) y 3D en una misma ventana, comenzando en pausa.
    """
    # Configuración de la figura
    fig = plt.figure(figsize=(16, 7))
    
    # ---- Subplot 1: Animación 2D (XY) ----
    ax_2d = fig.add_subplot(121)
    ax_2d.set_xlim(min(x)-0.1, max(x)+0.1)
    ax_2d.set_ylim(min(y)-0.1, max(y)+0.1)
    ax_2d.set_xlabel('$x$', fontsize=12)
    ax_2d.set_ylabel('$y$', fontsize=12)
    title_2d = 'Trayectoria 2D (XY)'
    if wc is not None and wz is not None:
        title_2d += f'\n$\omega_c = {wc}$, $\omega_z = {wz}$'
    ax_2d.set_title(title_2d, fontsize=14)
    ax_2d.grid(True)

    line_2d, = ax_2d.plot([], [], 'b-', lw=1, alpha=0.6)
    point_2d, = ax_2d.plot([], [], 'ro', markersize=6)

    # ---- Subplot 2: Animación 3D ----
    ax_3d = fig.add_subplot(122, projection='3d')
    ax_3d.set_xlim(min(x), max(x))
    ax_3d.set_ylim(min(y), max(y))
    ax_3d.set_zlim(min(z), max(z))
    ax_3d.set_xlabel('$x$', fontsize=12)
    ax_3d.set_ylabel('$y$', fontsize=12)
    ax_3d.set_zlabel('$z$', fontsize=12)
    title_3d = 'Trayectoria 3D'
    if wc is not None and wz is not None:
        title_3d += f'\n$\omega_c = {wc}$, $\omega_z = {wz}$'
    ax_3d.set_title(title_3d, fontsize=14)

    line_3d, = ax_3d.plot([], [], [], 'b-', lw=1, alpha=0.6)
    point_3d, = ax_3d.plot([], [], [], 'ro', markersize=6)

    # ---- Botón Play/Pause ----
    ax_button = plt.axes([0.4, 0.02, 0.2, 0.05])  # Posición del botón [x, y, ancho, alto]
    button = Button(ax_button, 'Play', color='lightgoldenrodyellow', hovercolor='0.8')

    # Función de inicialización
    def init():
        line_2d.set_data([], [])
        point_2d.set_data([], [])
        line_3d.set_data([], [])
        line_3d.set_3d_properties([])
        point_3d.set_data([], [])
        point_3d.set_3d_properties([])
        return line_2d, point_2d, line_3d, point_3d

    # Función de actualización
    def update(frame):
        line_2d.set_data(x[:frame*skip_points], y[:frame*skip_points])
        point_2d.set_data(x[frame*skip_points], y[frame*skip_points])
        line_3d.set_data(x[:frame*skip_points], y[:frame*skip_points])
        line_3d.set_3d_properties(z[:frame*skip_points])
        point_3d.set_data(x[frame*skip_points], y[frame*skip_points])
        point_3d.set_3d_properties(z[frame*skip_points])
        ax_3d.view_init(elev=20, azim=frame/2)
        return line_2d, point_2d, line_3d, point_3d

    # Crea la animación (inicialmente pausada)
    frames = len(t) // skip_points
    ani = FuncAnimation(fig, update, frames=frames, init_func=init, 
                        blit=True, interval=30, repeat=True)
    ani.event_source.stop()  # Pausa la animación al inicio

    # Función para toggle Play/Pause
    def toggle_animation(event):
        if ani.event_source:
            ani.event_source.stop()
            button.label.set_text('Play')
        else:
            ani.event_source.start()
            button.label.set_text('Pause')
        fig.canvas.draw()

    button.on_clicked(toggle_animation)

    plt.tight_layout()
    plt.show()
    return ani

# --- Ejemplo de uso ---
if __name__ == "__main__":
    # Genera datos de ejemplo (simulación de trampa de Penning)
    t = np.linspace(0, 20, 2000)
    wc, wz = 1.0, 0.5
    x = np.cos(wc*t) * np.cos(wz*t)
    y = np.sin(wc*t) * np.cos(wz*t)
    z = np.sin(wz*t)
    
    # Llama a la función
    ani = animar_trayectorias_combinadas(x, y, z, t, wc=wc, wz=wz, skip_points=10)
    
    # Para guardar la animación (opcional)
    # ani.save('trayectoria_combinada.mp4', writer='ffmpeg', fps=30)
    
    plt.show()