# -*- coding: utf-8 -*-
"""
Created on Thu May 22 15:40:25 2025

@author: Martin Drechsler
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
from mpl_toolkits.mplot3d import Axes3D

from penningAni import animar_trayectorias_combinadas, animar_trayectorias_con_boton
#%%
# Parámetros del sistema (ajustables)
wc = 1 # Frecuencia ciclotrón (rad/s)
wz = .5  # Frecuencia axial (rad/s)
print(wc**2-2*wz**2)
#%%
# Condiciones iniciales [x, y, z, vx, vy, vz]
initial_conditions = [1.0, 0.0, 0.5, 0.0, 1.0, 0.0]  # x0, y0, z0, vx0, vy0, vz0

# Tiempo de simulación
t_span = (0, 200)  # Intervalo de tiempo (s)
t_eval = np.linspace(t_span[0], t_span[1], 3000)  # Puntos de evaluación

# Sistema de ecuaciones diferenciales
def equations(t, state, wc, wz):
    x, y, z, vx, vy, vz = state
    dxdt = vx
    dydt = vy
    dzdt = vz
    dvxdt = (wc**2) * vy + (wz**2 / 2) * x
    dvydt = -(wc**2) * vx + (wz**2 / 2) * y
    dvzdt = -wz**2 * z  # Ecuación para z
    return [dxdt, dydt, dzdt, dvxdt, dvydt, dvzdt]

# Solución numérica
solution = solve_ivp(equations, t_span, initial_conditions, args=(wc, wz),
                    t_eval=t_eval, method='RK45', rtol=1e-8)

# Extraer resultados
x = solution.y[0]
y = solution.y[1]
z = solution.y[2]
t = solution.t

# Gráfica de la trayectoria
plt.figure(figsize=(10, 6))
plt.plot(x, y, 'b-', linewidth=1, label='Trayectoria (x, y)')
plt.xlabel('Posición $x$', fontsize=12)
plt.ylabel('Posición $y$', fontsize=12)
plt.title(f'Trayectoria en una trampa de Penning\n$\omega_c = {wc}$, $\omega_z = {wz}$', fontsize=14)
plt.grid(True)
plt.axis('equal')
plt.legend()
plt.show()

# Gráfica 3D de la trayectoria
fig = plt.figure(figsize=(12, 8))
ax = fig.add_subplot(111, projection='3d')
ax.plot(x, y, z, 'b-', linewidth=0.8, alpha=0.8)
ax.set_xlabel('Posición $x$', fontsize=12)
ax.set_ylabel('Posición $y$', fontsize=12)
ax.set_zlabel('Posición $z$', fontsize=12)
ax.set_title(f'Trayectoria 3D en una trampa de Penning\n$\omega_c = {wc}$, $\omega_z = {wz}$', fontsize=14)
plt.tight_layout()

# Gráficas de posición vs tiempo
fig2, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(10, 8), sharex=True)
ax1.plot(t, x, 'r-', label='$x(t)$')
ax1.set_ylabel('$x$', fontsize=12)
ax1.legend()
ax1.grid(True)

ax2.plot(t, y, 'g-', label='$y(t)$')
ax2.set_ylabel('$y$', fontsize=12)
ax2.legend()
ax2.grid(True)

ax3.plot(t, z, 'b-', label='$z(t)$')
ax3.set_ylabel('$z$', fontsize=12)
ax3.set_xlabel('Tiempo $t$', fontsize=12)
ax3.legend()
ax3.grid(True)

fig2.suptitle('Evolución temporal de las coordenadas', fontsize=14)
plt.tight_layout()
plt.show()

#%%
# ani = animar_trayectorias_combinadas(x, y, z, t, wc=wc, wz=wz, skip_points=10)
ani = animar_trayectorias_con_boton(x, y, z, t, wc=wc, wz=wz, skip_points=10)
plt.show()