#QUADRATIC FRICTION - HEUN METHOD 
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
θ=60*pi/180 #angulo
vo=100      # velocidad m/s

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

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

# 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)
    i
    if y[i] < 0:  # Si el proyectil toca el suelo, detener la simulación
            break

# Gráfica de la trayectoria
plt.plot(x, y)
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()

# Imprimir el valor del tiempo
print("El valor del tienpo de vuelo es:", i*dt)