import numpy as np
import matplotlib.pyplot as plt

# Constantes do problema
d = 0.1        # diâmetro das esferas (m)
l = 10 * d     # comprimento das cordas (m)
m = 0.3        # massa (kg)
g = 9.8        # aceleração gravítica (m/s²)
k = 1e7        # constante elástica (N/m²)
q = 2.0        # expoente do toque


def acc_toque(dx,d):
    # calcular a aceleração de uma esfera devido ao contacto com a esfera à sua direita
    k = 1e7
    q = 2.0
    if dx<d:
        a = (-k*abs(dx-d)**q)/m
    else:
        a = 0.0
    return a

def acc_i(i,x,N):
    #calcular a aceleração de esfera i, cuja posicao de equilibrio é d*i
    a = 0

    if i>0: # a primeira esfera não tem vizinho à sua esquerda
        a -= acc_toque(x[i] - x[i-1],d)
    if i < (N-1): # a última esfera não tem vizinho à sua direita
        a += acc_toque(x[i+1] - x[i], d)

    # aceleração de gravidade, afeta todas as esferas
    a -= g*(x[i]-d*i)/l
    return a

def simular_esferas(N, x0, v0, dt=1e-4, t_max=5.0):
    """Simula o sistema de N esferas com Euler-Cromer."""
    n_passos = int(t_max / dt)
    x = np.zeros((n_passos, N))
    v = np.zeros((n_passos, N))
    a = np.zeros(N)

    # Condições iniciais
    x[0] = x0
    v[0] = v0

    # Simulação temporal
    for t in range(1, n_passos):
        for i in range(N):
            a[i] = acc_i(i, x[t - 1], N)
            v[t, i] = v[t - 1, i] + a[i] * dt
            x[t, i] = x[t - 1, i] + v[t, i] * dt

    tempo = np.linspace(0, t_max, n_passos)
    return tempo, x, v

def plot_posicoes(tempo, x, titulo='Posição das esferas'):
    for i in range(x.shape[1]):
        plt.plot(tempo, x[:, i], label=f'Esfera {i}')
    plt.xlabel("Tempo (s)")
    plt.ylabel("Posição x (m)")
    plt.title(titulo)
    plt.legend()
    plt.grid(True)
    plt.show()


#1 Simule o movimento de duas esferas sujeitas às forças descritas na pagina anterior
# Simule o movimento durante 5 segundos.

N = 2
x0 = np.array([-5*d, d])  # posições iniciais
v0 = np.zeros(N)          # velocidades iniciais
tempo, x, v = simular_esferas(N, x0, v0)
plot_posicoes(tempo, x, "Exercício 1 - Duas esferas")



## 3 Calcule a energia cinética e a energia potencial em cada instante

def calcular_energia_momento(x, v, d=d, l=l, m=m):
    Ep = 0.5 * m * g / l * (x - d * np.arange(x.shape[1]))**2
    Ec = 0.5 * m * v**2
    energia_total = np.sum(Ep + Ec, axis=1)
    momento_total = np.sum(m * v, axis=1)
    return energia_total, momento_total

def plot_energia_momento(tempo, energia, momento):
    plt.plot(tempo, energia, label="Energia Total")
    plt.xlabel("Tempo (s)")
    plt.ylabel("Energia (J)")
    plt.title("Energia total no sistema")
    plt.grid()
    plt.show()

    plt.plot(tempo, momento, label="Momento Linear Total")
    plt.xlabel("Tempo (s)")
    plt.ylabel("Momento (kg·m/s)")
    plt.title("Momento linear total do sistema")
    plt.grid()
    plt.show()


energia, momento = calcular_energia_momento(x, v)
plot_energia_momento(tempo, energia, momento)


#########################
# Exercicio 2 : Múltiplas esferas
#  Adapte o código do exercício anterior para simular o movimento de 𝑁 esferas.
# Cada esfera elevada deve começar a uma distância de −5𝑑 da sua posição de equilíbrio.

N = 5
x0 = np.array([-5*d, -4*d] + [i*d for i in range(2, N)])
v0 = np.zeros(N)
tempo, x, v = simular_esferas(N, x0, v0)
plot_posicoes(tempo, x, "Exercício 2 - 5 esferas com 2 levantadas")