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

#Exercício 2: Integração numerérica
#1. Calcule a energia mecânica de 𝑡0 = 0 até o momento em que a bola bate no solo ( 𝑡𝑓 ≈ 0.4 s).
#A massa da bola é 57 g.
#(dados do exercício 1)
pos = (0,2,3) #m
x0,y0,z0 = pos

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

g = 9.8
vT = 120000/3600 #m/s
dt = 0.01
t0 = 0.0
tf = 0.4    # bola bate no solo (ex2)
m = 0.057   #massa da bola (kg)
n = int((tf - t0) / dt)

D = g/vT**2

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) 

    ax[i] = -D * vv * vx[i]     
    ay[i] = -D * vv * vy[i]    
    az[i] = -g - D * vv * vz[i]  

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

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


    if z[i+1] <= 0:
        x = x[:i+2]
        y = y[:i+2]
        z = z[:i+2]
        t = t[:i+2]

        vx = vx[:i+2]
        vy = vy[:i+2]
        vz = vz[:i+2]

        break

#energia mecanica:
# Energia cinética: (1/2) m (vx² + vy² + vz²)
# Energia potencial: m g z
E = 0.5 * m * (vx**2 + vy**2 + vz**2) + m * g * z

print("Energia mecânica inicial: {:.3f} J".format(E[0]))
print("Energia mecânica no impacto: {:.3f} J".format(E[-1]))

plt.figure(figsize=(8,6))
plt.plot(t, E, 'r-', label='Energia mecânica')
plt.xlabel('Tempo (s)')
plt.ylabel('Energia Mecânica (J)')
plt.title('Evolução da Energia Mecânica da bola')
plt.legend()
plt.grid(True)
plt.show()