#To run the program with greater precision, for example with dt=0.0001, 
#it is recommended to run this file in JupyterLab.
import numpy as np
import matplotlib.pyplot as plt
from math import cos, sin, pi

# Parámetros
m = 1     # masa kg
b = 0.1   # coeficiente de arrastre
g = 9.81  # gravedad
vo = 700  # velocidad inicial m/s

# Condiciones iniciales
x0 = 0
y0 = 0

# Tiempo
t0 = 0
tf = 1000
dt = 0.001  # paso de tiempo

# Función para calcular aceleraciones
def ac(vx, vy, m, b, g):
    ax = -b*vx*np.sqrt(vx**2 + vy**2) / m
    ay = -g - b*vy*np.sqrt(vx**2 + vy**2) / m
    return ax, ay

# Método de Runge-Kutta para resolver las ecuaciones
def rk4(m, b, g, te, vo, dt):
    t = np.arange(t0, tf, dt)
    
    # Arrays para almacenar soluciones
    x = np.zeros_like(t)
    y = np.zeros_like(t)
    vx = np.zeros_like(t)
    vy = np.zeros_like(t)
    U = np.zeros_like(t)  # Energía potencial

    # Condiciones iniciales
    vx[0] = vo*cos(te)  
    vy[0] = vo*sin(te)

    # Variables auxiliares
    tv = 0  # tiempo de vuelo
    ymax = 0  # altura máxima
    tm = 0  # tiempo en altura máxima

    for i in range(1, len(t)):
        ax1, ay1 = ac(vx[i-1], vy[i-1], m, b, g)
        k1vx = ax1*dt
        k1vy = ay1*dt
        k1x = vx[i-1]*dt
        k1y = vy[i-1]*dt

        ax2, ay2 = ac(vx[i-1] + 0.5*k1vx, vy[i-1] + 0.5*k1vy, m, b, g)
        k2vx = ax2*dt
        k2vy = ay2*dt
        k2x = (vx[i-1] + 0.5*k1vx)*dt
        k2y = (vy[i-1] + 0.5*k1vy)*dt

        ax3, ay3 = ac(vx[i-1] + 0.5*k2vx, vy[i-1] + 0.5*k2vy, m, b, g)
        k3vx = ax3*dt
        k3vy = ay3*dt
        k3x = (vx[i-1] + 0.5*k2vx)*dt
        k3y = (vy[i-1] + 0.5*k2vy)*dt

        ax4, ay4 = ac(vx[i-1] + k3vx, vy[i-1] + k3vy, m, b, g)
        k4vx = ax4*dt
        k4vy = ay4*dt
        k4x = (vx[i-1] + k3vx)*dt
        k4y = (vy[i-1] + k3vy)*dt

        # Actualizar velocidades y posiciones
        vx[i] = vx[i-1] + (k1vx + 2*k2vx + 2*k3vx + k4vx)/6
        vy[i] = vy[i-1] + (k1vy + 2*k2vy + 2*k3vy + k4vy)/6
        x[i] = x[i-1] + (k1x + 2*k2x + 2*k3x + k4x)/6
        y[i] = y[i-1] + (k1y + 2*k2y + 2*k3y + k4y)/6

        # Energía potencial
        U[i] = m*g*y[i]

        # Verificar si ha alcanzado el suelo
        if y[i] <= 0:
            tv = t[i]  # tiempo de vuelo
            break

    return x[:i+1], y[:i+1], U[:i+1], t[:i+1], tv

# Ángulos de lanzamiento en radianes
angles = [30*pi/180, 45*pi/180, 60*pi/180]
labels = ['30°', '45°', '60°']
colors = ['blue', 'green', 'red']

# Graficar la energía potencial para cada ángulo
plt.figure(figsize=(10, 6))
for te, label, color in zip(angles, labels, colors):
    x, y, U, t, tv = rk4(m, b, g, te, vo, dt)

    # Filtrar los datos para el intervalo [0, tv]
    t_filtered = t[(t >= 0) & (t <= tv)]
    U_filtered = U[(t >= 0) & (t <= tv)]

    # Graficar energía potencial
    plt.plot(t_filtered, U_filtered, label=f'{label}', color=color)

# Personalización del gráfico
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$U(J)$', fontsize=12)
plt.xticks(fontsize=12)
plt.yticks(np.arange(0, np.max(U_filtered)+10, 100), fontsize=12)
plt.grid(True)

# Leyenda
legend = plt.legend(facecolor='white', framealpha=1, edgecolor='black', fontsize=12, loc='upper right')
legend.get_frame().set_linewidth(1)
plt.tight_layout()
plt.show()