import numpy as np
import matplotlib.pyplot as plt

# Constantes do problema
m = 2000           # massa (kg)
theta = 5          # inclinação (graus)
P_subida = 40000   # W (potência positiva na subida)
P_descida = -15000 # W (potência negativa na descida)
v0_subida = 1      # m/s
v0_descida = 20    # m/s
d_total = 2000     # metros

mu = 0.04
Cres = 0.25
A = 2
rho_ar = 1.225


# ========= Funções gerais =========

def forca_gravidade(m, g, theta):
    return -m * g * np.sin(np.radians(theta))

def forca_rolamento(m, g, mu, theta):
    return -mu * m * g * np.cos(np.radians(theta))

def forca_ar(Cres, A, rho_ar, v):
    return -0.5 * Cres * A * rho_ar * abs(v) * v

def forca_motor(P, v):
    return P / v

def simular_movimento(m, P, v0, d_objetivo, theta, mu, Cres, A, rho_ar, g=9.8, dt=0.1):
    """
    Simula o movimento do carro usando Euler explícito.
    
    Retorna: tempo_total, distancia_total, trabalho_motor, arrays de tempo, posição, velocidade
    """
    t = [0]
    v = [v0]
    x = [0]
    W_motor = [0]

    i = 0
    while x[-1] < d_objetivo:
        fg = forca_gravidade(m, g, theta)
        fr = forca_rolamento(m, g, mu, theta)
        fa = forca_ar(Cres, A, rho_ar, v[i])
        fm = forca_motor(P, v[i])
        
        if abs(v[i]) < 0.5:  # carro quase parando
            print("Carro parou ou quase parou. Interrompendo simulação.")
            break

        F_total = fg + fr + fa + fm
        a = F_total / m
        
        v_new = v[i] + a * dt
        x_new = x[i] + v[i] * dt
        t_new = t[i] + dt
        W_new = W_motor[i] + P * dt  # potência constante

        # Append
        v.append(v_new)
        x.append(x_new)
        t.append(t_new)
        W_motor.append(W_new)
        
        i += 1

    return t[-1], x[-1], W_motor[-1], t, x, v


#1 evolução temporal da posição e da velocidade do carro, 
# se o carro produzir continuamente a potência 40kW, e partir de uma velocidade de 1m/s.

# Simulação da subida
t_subida, x_subida, W_subida, t1, x1, v1 = simular_movimento(m, P_subida, v0_subida, d_total, theta, mu, Cres, A, rho_ar)

# Plotar evolução temporal da posição e velocidade
fig, ax = plt.subplots(2, 1, figsize=(10, 8))

ax[0].plot(t1, x1, label="Posição (m)")
ax[0].set_ylabel("Posição (m)")
ax[0].grid(True)
ax[0].legend()

ax[1].plot(t1, v1, label="Velocidade (m/s)", color='orange')
ax[1].set_xlabel("Tempo (s)")
ax[1].set_ylabel("Velocidade (m/s)")
ax[1].grid(True)
ax[1].legend()

plt.suptitle("Evolução da posição e velocidade do carro na subida")
plt.tight_layout()
plt.show()
 
print("Tempo total de subida (2km): {:.2f} s".format(t_subida))
print(f"Trabalho do motor: {W_subida/1000:.2f} kJ")


########################################################
#4 Numa segunda fase, o carro agora desce a mesma inclinação, começando com velocidade 20m/s. 
# Usando travagem regenerativa, o carro aplica continuamente uma potência -15kW.

t_descida, x_descida, W_descida, t2, x2, v2 = simular_movimento(m, P_descida, v0_descida, d_total, theta, mu, Cres, A, rho_ar)

print(f"Descida - Tempo: {t_descida:.1f} s")
print(f"Trabalho do motor: {W_descida/1000:.2f} kJ")

#5. Se 50% do trabalho na descida é recuperado para carregar a bateria do carro, 
#qual a diferência de energia na bateria no final, depois de ter feito a subida e a descida, 
# comparado com no início? (Assume 100% eficiencia do motor na subida.)

# Recuperação de 50% do trabalho da descida
recuperado = 0.5 * abs(W_descida)

energia_total =  W_subida + recuperado

print(f"Diferença final de energia na bateria: {energia_total/1000:.2f} kJ")


plt.figure(figsize=(10, 5))

plt.subplot(1, 2, 1)
plt.plot(t1, x1, label='Subida')
plt.plot(t2, x2, label='Descida')
plt.xlabel("Tempo (s)")
plt.ylabel("Posição (m)")
plt.title("Evolução da posição")
plt.legend()

plt.subplot(1, 2, 2)
plt.plot(t1, v1, label='Subida')
plt.plot(t2, v2, label='Descida')
plt.xlabel("Tempo (s)")
plt.ylabel("Velocidade (m/s)")
plt.title("Evolução da velocidade")
plt.legend()

plt.tight_layout()
plt.grid(True)
plt.show()