import numpy as np
import matplotlib.pyplot as plt

# ======================================================
# Dados do Exercício 3D (adaptado dos dados do exercício 1)
# ======================================================
# Posição inicial (x0, y0, z0) em metros
x0, y0, z0 = 0, 2, 3

# Velocidades iniciais (convertidas de km/h para m/s)
vx0 = 160000/3600   # aproximadamente 44.44 m/s
vy0 = 20000/3600    # aproximadamente 5.56 m/s
vz0 = -20000/3600   # aproximadamente -5.56 m/s

g = 9.8             # aceleração da gravidade (m/s²)
vT = 120000/3600    # velocidade terminal (m/s)
dt = 0.001           # passo de integração (s)
t0 = 0.0
tf = 0.4          # tempo final (aproximadamente o instante de impacto)
m = 0.057           # massa da bola (kg)

# Parâmetro de arrasto (drag)  
# para a aceleração de arrasto: a_res = - d * |v| * v.
d = g/(vT**2)

n = int((tf - t0)/dt) + 1

t = np.empty(n)
x = np.empty(n)
y = np.empty(n)
z = np.empty(n)

vx = np.empty(n)
vy = np.empty(n)
vz = np.empty(n)

ax = np.empty(n)
ay = np.empty(n)
az = np.empty(n)

# Aceleração causada pela resistência (para cada componente)
aresx = np.empty(n)
aresy = np.empty(n)
aresz = np.empty(n)

# Outras variáveis:
v = np.empty(n)          # módulo da velocidade
Em = np.empty(n)         # Energia mecânica (cinética + potencial)
fun = np.empty(n)        # integrando para o trabalho da resistência
int_trab_res = np.empty(n)  # trabalho acumulado (integral)

# Condições iniciais
t[0] = t0
x[0], y[0], z[0] = x0, y0, z0
vx[0], vy[0], vz[0] = vx0, vy0, vz0

# No instante inicial o trabalho acumulado é zero
int_trab_res[0] = 0.0

# ========================================================
# Integração Numérica (método de Euler para movimento 3D)
# ========================================================

for i in range(n-1):
    # Calcula o módulo da velocidade no passo i
    v[i] = np.sqrt(vx[i]**2 + vy[i]**2 + vz[i]**2)
    
    # Calcula as acelerações de arrasto para cada componente
    aresx[i] = -d * v[i] * vx[i]
    aresy[i] = -d * v[i] * vy[i]
    aresz[i] = -d * v[i] * vz[i]
    
    # Aceleração total: 
    # Para x e y só temos a resistência; para z, somamos a gravidade (que atua para baixo)
    ax[i] = aresx[i]
    ay[i] = aresy[i]
    az[i] = aresz[i] - g
    
    # Atualiza as velocidades
    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 as posições
    x[i+1] = x[i] + vx[i] * dt
    y[i+1] = y[i] + vy[i] * dt
    z[i+1] = z[i] + vz[i] * dt
    
    # Calcula a energia mecânica em cada instante (Ec + Ep)
    Em[i] = 0.5 * m * v[i]**2 + m * g * z[i]
    
    # Calcula o integrando para o trabalho da resistência:
    # fun = m * (aresx * vx + aresy * vy + aresz * vz)
    fun[i] = m * ( aresx[i]*vx[i] + aresy[i]*vy[i] + aresz[i]*vz[i] )
    
    int_trab_res[i] = dt * ( 0.5*(fun[0] + fun[i]) + np.sum(fun[1:i]) )

    
    # Verifica se a bola atingiu o solo (z <= 0) para interromper a simulação.
    if z[i+1] <= 0:
        # Trunca os arrays até o instante do impacto
        t = t[:i+2]
        x = x[:i+2]
        y = y[:i+2]
        z = z[:i+2]
        vx = vx[:i+2]
        vy = vy[:i+2]
        vz = vz[:i+2]
        Em = Em[:i+2]
        fun = fun[:i+2]
        int_trab_res = int_trab_res[:i+2]
        break

v[len(t)-1] = np.sqrt(vx[-1]**2 + vy[-1]**2 + vz[-1]**2)
Em[len(t)-1] = 0.5 * m * v[-1]**2 + m * g * z[-1]

# =====================================================
# Queremos o trabalho acumulado em: t = 0, t = 0.2 e t = 0.4 s.
indices = [int(round(target/dt)) for target in [0.0, 0.2, 0.4]]
print(int_trab_res)
print(indices)
print("Energia mecânica inicial: {:.3f} J".format(Em[0]))
print("Energia mecânica no impacto: {:.3f} J".format(Em[-1]))


print("\nTrabalho acumulado da resistência do ar:")
print("De t = 0.0 s: {:.4f} J".format(int_trab_res[indices[0]]))
print("De t = 0.2 s: {:.4f} J".format(int_trab_res[indices[1]]))
print("De t = 0.4 s: {:.4f} J".format(int_trab_res[indices[2]-1]))

if (Em[-1] - Em[0]) + int_trab_res[[indices[2]-1]] == 0:
    print("Conservação da energia confirmada.")
else:
    print("Conservação da energia não confirmada, erro de {} J".format((Em[-1] - Em[0] - int_trab_res[[indices[2]-1]])))


# =====================================================
# Gráficos: Trajetória 3D e trabalho acumulado vs tempo
# =====================================================
from mpl_toolkits.mplot3d import Axes3D  # para plotagem 3D

fig = plt.figure(figsize=(10,6))
ax3d = fig.add_subplot(121, projection='3d')
ax3d.plot(x, y, z, 'b-', label='Trajetória')
ax3d.set_xlabel('x (m)')
ax3d.set_ylabel('y (m)')
ax3d.set_zlabel('z (m)')
ax3d.set_title('Trajetória da bola 3D')
ax3d.legend()

plt.subplot(122)
plt.plot(t, int_trab_res, 'r-', label='Trabalho acumulado')
plt.xlabel('Tempo (s)')
plt.ylabel('Trabalho (J)')
plt.title('Trabalho realizado pela resistência')
plt.legend()
plt.grid(True)
plt.tight_layout()
plt.show()


#####################################################
#3  Calcule o trabalho realizado pela força de resistência do ar usando a conservação de energia
W_res_energia = Em[-1] - Em[0]
print("\nTrabalho pela resistência (por conservação da energia): {:.4f} J".format(W_res_energia))