#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
t = np.arange(t0, tf, dt)

# 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 Runge-Kutta de orden 4 con cálculo de máximos
def rk4(m, b, g, te, vo, dt, t):
    # Arrays para almacenar soluciones
    x = np.zeros_like(t)
    y = np.zeros_like(t)
    vx = np.zeros_like(t)
    vy = np.zeros_like(t)

    # Condiciones iniciales
    vx0 = vo*cos(te)  # velocidad inicial en x
    vy0 = vo*sin(te)  # velocidad inicial en y
    x[0] = x0
    y[0] = y0
    vx[0] = vx0
    vy[0] = vy0

    ymax = 0
    xmax = 0
    xm = 0
    tm = 0
    tv = 0

    # Resolver las ecuaciones con el método de Runge-Kutta
    for i in range(1, len(t)):
        # Incremente intermedios para vx y vy
        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

        # Comprobar si alcanzó la altura máxima
        if vy[i-1] > 0 and vy[i] <= 0:
            tm = t[i]
            xm = x[i]
            ymax = y[i]

        # Comprobar si llegó al suelo
        if y[i] <= 0:
            tv = t[i]
            xmax = x[i]
            break

    # Retornar las soluciones y los puntos máximos
    return x[:i+1], y[:i+1], xm, ymax, xmax, tm, tv

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

# Graficar las trayectorias para cada ángulo
plt.figure(figsize=(8,6))
for te, label, color in zip(angles, labels, colors):
    x1, y1, xm, ymax, xmax, tm, tv = rk4(m, b, g, te, vo, dt, t)
    plt.plot(x1, y1, label=f'{label}', color=color)
    plt.scatter(xm, ymax, color=color, s=25, marker='o')  # Punto de altura máxima

# Personalización del gráfico
plt.xlabel('$x(m)$', fontsize=12)
plt.ylabel('$y(m)$', fontsize=12)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.grid(True)

# Establecer leyenda con fondo blanco y sin transparencia
legend = plt.legend(facecolor='white', framealpha=1, edgecolor='black', fontsize=12)
legend.get_frame().set_linewidth(1)  # Grosor del borde de la leyenda

plt.axis('equal')
plt.tight_layout()
plt.show()