#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 constantes
m = 1     # masa en kg
b = 0.1   # coeficiente de arrastre
g = 9.81  # aceleración debida a la 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)
    Et = np.zeros_like(t)

    # Condiciones iniciales
    vx[0] = vo*cos(te)  
    vy[0] = vo*sin(te)
    Et[0] = m*g*y[0] + m*(vx[0]**2 + vy[0]**2) / 2  # Energía cinética inicial

    # 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 cinética
        Et[i] = m*g*y[i] + m*(vx[i]**2 + vy[i]**2)/2

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

    return x[:i+1], y[:i+1], Et[: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 cinética para cada ángulo
plt.figure(figsize=(10, 6))
for te, label, color in zip(angles, labels, colors):
    x, y, Et, t, tv = rk4(m, b, g, te, vo, dt)

    # Filtrar los datos para el intervalo [0.5, tv]
    t_filtered = t[(t >= 0.5) & (t <= tv)]
    Et_filtered = Et[(t >= 0.5) & (t <= tv)]

    # Graficar energía cinética
    plt.plot(t_filtered, Et_filtered, label=f'{label}', color=color)

# Personalización del gráfico
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$E(J)$', fontsize=12)
plt.xticks(fontsize=12)
plt.yticks(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()