#QUADRATIC FRICTION - EULER METHOD 
import numpy as np
import matplotlib.pyplot as plt
from math import cos, sin, pi

# Parámetros
m=1  # masa kg
b=0.001  # coeficiente de arrastre
g=9.81  # aceleración debido a la gravedad

# Condiciones iniciales
θ=60*pi/180 #angulo
vo=100 # velocidad m/s
x0=0
y0=0
vx0=vo*cos(θ) # velocidad inicial en x
vy0=vo*sin(θ) # velocidad inicial en y

# Tiempo
t0=0.0
tf=1000
dt=0.001  # paso del tiempo
t=np.arange(t0,tf,dt)

# Arrays para almacenar las 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)
K=np.zeros_like(t) 

# Semillas
x[0]=x0
y[0]=y0
vx[0]=vx0     # velocidad inicial en x
vy[0]=vy0     # velocidad inicial en y
K[0]=m*(vx0**2+vy0**2)/2  # velocidad total inicial

# Resolver las ecuaciones usando el método de Euler
for i in range(1,len(t)):
    # Calcular aceleraciones
    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
    
    # Actualizar velocidades
    vx[i]=vx[i-1]+ax*dt
    vy[i]=vy[i-1]+ay*dt
    
    # Actualizar posiciones
    x[i]=x[i-1]+vx[i]*dt
    y[i]=y[i-1]+vy[i]*dt
    
    # Calculo del cuadrado de la velocidad
    K[i]=m*(vx[i]**2+vy[i]**2)/2

    # Calculo del radio vector
    r[i]=np.sqrt(x[i]**2+y[i]**2)
    if y[i] < 0:  # Si el proyectil toca el suelo, detener la simulación
        break

# Graficar solo los puntos válidos
x_valid=x[:i]
y_valid=y[:i]
r_valid=r[:i]
vx_valid=vx[:i]
vy_valid=vy[:i]
t_valid=t[:i]

# Gráfica de la trayectoria
plt.figure(figsize=(8,6))
plt.plot(x_valid,y_valid)
plt.xlabel('$x(m)$',fontsize=16)
plt.ylabel('$y(m)$',fontsize=16)
plt.grid(True)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.ylim(0, 300)       # Aquí fijas el rango de y: de 0 hasta a
plt.axis('auto')     # Opcional si no quieres 'equal'
plt.tight_layout()
plt.show()

# Energía cinética en función del tiempo
plt.figure(figsize=(7,6))
plt.plot(t[:i],K[:i])
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$K(J)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Espacio de fase de x
plt.figure(figsize=(7,6))
plt.plot(x_valid,m*vx_valid)
plt.xlabel('$x$', fontsize=12)
plt.ylabel('$p_x$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Espacio de fase de y
plt.figure(figsize=(7,6))
plt.plot(y_valid,m*vy_valid)
plt.xlabel('$y$', fontsize=12)
plt.ylabel('$p_y$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show()

# Radio vector en función del tienpo
plt.figure(figsize=(7,6))
plt.plot(t[:i],r[:i])
plt.xlabel('$t(s)$', fontsize=12)
plt.ylabel('$r(m)$', fontsize=12)
plt.grid(True)
plt.xticks(fontsize=12)
plt.yticks(fontsize=12)
plt.tight_layout()
plt.show() 

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