#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, sqrt

# 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

# 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)
    r = np.zeros_like(t)

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

    # Variables auxiliares
    tv = 0  # tiempo de vuelo

    for i in range(1, len(t)):
        # Método de Runge-Kutta (RK4)
        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

        # Calcular el radio vector
        r[i] = sqrt(x[i]**2 + y[i]**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], r[: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 el radio vector para cada ángulo
plt.figure(figsize=(10, 6))
for te, label, color in zip(angles, labels, colors):
    x, y, r, t, tv = rk4(m, b, g, te, vo, dt)

    # Filtrar los datos para el intervalo [0, tv]
    t_filtered = t[t <= tv]
    r_filtered = r[t <= tv]

    # Graficar radio vector
    plt.plot(t_filtered, r_filtered, label=f'{label}', color=color)

# Personalización del gráfico
plt.xlabel('$t(s)$', fontsize=16)
plt.ylabel('$r(m)$', fontsize=16)
plt.xticks(fontsize=16)
plt.yticks(fontsize=16)
plt.grid(True)

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