#QUADRATIC FRICTION - HEUN METHOD (Security parable)
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

# Tiempo
t0=0
tf=10
dt=0.001  # paso de tiempo
t=np.arange(t0,tf,dt)

# Bucle para ángulos de 0 a 90° en incrementos de 5°
for angulo in range(0,91,5):
    θ=angulo*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
    
    # Array para almacenar las soluciones
    x=np.zeros_like(t)
    y=np.zeros_like(t)
    vx=np.zeros_like(t)
    vy=np.zeros_like(t)

    # 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
            break

    # Gráfica de la trayectoria
    plt.plot(x[:i], y[:i],color='blue')

# Configuración del gráfico
plt.xlabel("$x(m)$", fontsize=12)
plt.ylabel("$y(m)$", fontsize=12)
plt.grid(True)

# Ajustar el tamaño de la letra de los ticks
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.show()