import sympy as sy
import numpy as np
import matplotlib.pyplot as plt

# Exercício 1. Servico de Tenis

pos = (0,2,3) #m
x0,y0,z0 = pos

v = (160000/3600,20000/3600,-20000/3600) #m/s
vx0,vy0,vz0 = v


#1  movimento da bola usando o método de Euler a 3D
#força de gravidade e a resistência do ar, com velocidade terminal 𝑣𝑇 = 120km/h.

g = 9.8
vT = 120000/3600 #m/s
dt = 0.01
t0 = 0.0
tf = 2.0
n = int((tf - t0) / dt)

D = g/vT**2 #coeficiente de arrasto

t = np.zeros(n+1)

x = np.zeros(n+1)
y = np.zeros(n+1)
z = np.zeros(n+1)
x[0] = x0  
y[0] = y0  
z[0] = z0  


vx = np.zeros(n+1)
vy = np.zeros(n+1)
vz = np.zeros(n+1)
vx[0] = vx0
vy[0] = vy0
vz[0] = vz0

ax = np.zeros(n+1)
ay = np.zeros(n+1)
az = np.zeros(n+1)


#euler
for i in range(n):
    t[i+1]=t[i]+dt
    vv=np.sqrt(vx[i]**2 + vy[i]**2 + vz[i]**2) #norma de v

    ax[i] = -D * vv * vx[i]       #resistencia do ar proporcional a velocidade
    ay[i] = -D * vv * vy[i]     #resistencia do ar proporcional a velocidade
    az[i] = -g - D * vv * vz[i]  #adiciona a gravidade

    #atualiza v
    vx[i+1] = vx[i] + ax[i] * dt
    vy[i+1] = vy[i] + ay[i] * dt
    vz[i+1] = vz[i] + az[i] * dt

    #atualiza pos
    x[i+1] = x[i] + vx[i] * dt
    y[i+1] = y[i] + vy[i] * dt
    z[i+1] = z[i] + vz[i] * dt

    #ex2  Em que ponto a bola cai no solo? 
    # -- Verifica se a bola atingiu o solo (z <= 0) --
    if z[i+1] <= 0:
        print(f"A bola atinge o solo em t = {t[i+1]:.2f} s")

        #retirar os 0s da lista
        x = x[:i+2]
        y = y[:i+2]
        z = z[:i+2]
        t = t[:i+2]

        break


fig = plt.figure(figsize=(10, 6))
ax = plt.axes(projection='3d') #criar eixos 3D
ax.plot(x, y, z, label='Trajetoria 3D')
ax.set_xlabel('x (m)')
ax.set_ylabel('y (m)')
ax.set_zlabel('z (m)')
ax.set_title('Simulacao da trajetoria usando o metodo de Euler 3D')
ax.legend()
plt.show()