#Time of flight as a function of angle - Heun 
import numpy as np
import matplotlib.pyplot as plt
from math import cos, sin, pi

# Parámetros
m=1         # masa
b=0.1       # coeficiente de fricción
g=9.81      # gravedad
vo=100      # velocidad m/s

# variación del tiempo y del ángulo inicial
t0=0
tf=10
dt=0.001  # paso de tiempo
t=np.arange(t0,tf,dt)
ꞵ=np.arange(0,91,5)

# Arrays para almacenar los resultados 
tv=[] #tiempo de vuelo
x=np.zeros_like(t)
y=np.zeros_like(t)
vx=np.zeros_like(t)
vy=np.zeros_like(t)

# Bucle para ángulos de 0 a 90° en incrementos de 5°
for an in ꞵ:
    θ=an*pi/180  # convertir el ángulo a radianes

    # Condiciones iniciales
    x0=0
    y0=0
    vx0=vo*cos(θ)  # velocidad inicial en x
    vy0=vo*sin(θ)  # velocidad inicial en y

    # Semilla
    x[0]=x0
    y[0]=y0
    vx[0]=vx0
    vy[0]=vy0

    # Resolver las ecuaciones usando el método de Heun
    for i in range(1,len(t)):
        # Aceleraciones actuales
        ax=-b*vx[i-1]*np.sqrt(vx[i-1]**2+vy[i-1]**2)/m
        ay=-g-b*vy[i-1]*np.sqrt(vx[i-1]**2+vy[i-1]**2)/m

        # Predicción con el método de Euler
        vxp=vx[i-1]+ax*dt
        vyp=vy[i-1]+ay*dt
        axp=-b*vxp*np.sqrt(vxp**2+vyp**2)/m
        ayp=-g-b*vyp*np.sqrt(vxp**2+vyp**2)/m

        # Corrección
        vx[i]=vx[i-1]+(dt/2)*(ax+axp)
        vy[i]=vy[i-1]+(dt/2)*(ay+ayp)
        x[i]=x[i-1]+(dt/2)*(vx[i]+vxp)
        y[i]=y[i-1]+(dt/2)*(vy[i]+vyp)

        if y[i] < 0:  # Si el proyectil toca el suelo, detener la simulación
            tv.append(i*dt)
            break

# Gráfica de ángulo contra tiempo final
plt.plot(ꞵ,tv)
plt.xlabel("Ángulo (θ°)",fontsize=12)
plt.ylabel("Tiempo final (s)",fontsize=12)
plt.grid(True)
plt.xticks(range(0, 91, 10),fontsize=12)
plt.yticks(fontsize=12)
plt.show()