import numpy as np
import matplotlib.pyplot as plt

#Exercício 1: Oscilador Acoplado
#2 corpos A e B acoplados através de uma mola de constante elástica k’
#Ligados a um ponto fixo através de molas de constante elástica k.

m = 1 #kg
kl = 0.5  #N/m
k= 1  #N/m
xA_eq = 1.0 #m
xB_eq = 2.0 #m

#i) 
xA0 = xA_eq + 0.3
xB0 = xB_eq + 0.3
vA0 = 0.0
vB0 = 0.0


def acelera(xi, m, k, kl, xA_eq, xB_eq):
    # u  = xA - xA_eq ;  up = du/dt
    # v  = xB - xB_eq ;  vp = dv/dt
    u = xi[0] - xA_eq
    v = xi[1] - xB_eq

    # equações de aceleração 
    #aA = ( Fk = -ku  )  + ( Fkp = -kl(u-v) )      /m  (f=ma)
    #aB = ( Fk = -kv  )  + ( Fkp = -kl(v-u) )    /m    (f=ma)
    aA = (-(k + kl)*u + kl*v) / m
    aB = ( kl*u     - (k + kl)*v) / m
    return np.array([aA, aB])

def rk4_x_vx(x, vx, acelera, m, k, kl, xA_eq, xB_eq, dt):
    ax1 = acelera(x, m, k, kl, xA_eq, xB_eq)
    c1v = ax1*dt
    c1x = vx*dt

    ax2 = acelera(x + c1x/2, m, k, kl, xA_eq, xB_eq)
    c2v = ax2*dt
    c2x = (vx + c1v/2.)*dt

    ax3 = acelera(x + c2x/2, m, k, kl, xA_eq, xB_eq)
    c3v = ax3*dt
    c3x = (vx + c2v/2.)*dt

    ax4 = acelera(x + c3x, m, k, kl, xA_eq, xB_eq)
    c4v = ax4*dt
    c4x = (vx + c3v)*dt

    xp = x + (c1x + 2.*c2x + 2.*c3x + c4x)/6.
    vxp = vx + (c1v + 2.*c2v + 2.*c3v + c4v)/6.
    return xp, vxp


def movimento_acoplado_2corpos_rk4(xA0, xB0, vA0, vB0, t_max=100):

    dt = 0.01
    n_passos = int(t_max / dt)
    t  = np.linspace(0, t_max, n_passos)
    x  = np.zeros((n_passos, 2))
    vx = np.zeros((n_passos, 2))
    x[0]  = [xA0,   xB0]
    vx[0] = [vA0,   vB0]


    #integrar as equações de movimento por rk4
    for i in range(1, n_passos):
        x[i], vx[i] = rk4_x_vx(x[i-1], vx[i-1],
            acelera, m, k, kl, xA_eq, xB_eq,
            dt)

    return t, x[:,0], x[:,1]

t, xA_t, xB_t = movimento_acoplado_2corpos_rk4(xA0, xB0, vA0, vB0)

plt.figure(figsize=(10, 5))
plt.plot(t, xA_t, label='xA(t)', color='blue')
plt.plot(t, xB_t, label='xB(t)', color='orange')
plt.xlabel('Tempo (s)')
plt.ylabel('Posição (m)')
plt.title('Movimento de dois corpos acoplados')
plt.legend()
plt.grid()
plt.show()

#ii)
xA0 = xA_eq + 0.3
xB0 = xB_eq - 0.3
vA0 = 0.0
vB0 = 0.0

t, xA_t, xB_t = movimento_acoplado_2corpos_rk4(xA0, xB0, vA0, vB0)
plt.figure(figsize=(10, 5))
plt.plot(t, xA_t, label='xA(t)', color='blue')
plt.plot(t, xB_t, label='xB(t)', color='orange')
plt.xlabel('Tempo (s)')
plt.ylabel('Posição (m)')
plt.title('Movimento de dois corpos acoplados (caso ii)')
plt.legend()
plt.grid()
plt.show()

#iii)
xA0 = xA_eq + 0.3
xB0 = xB_eq - 0.1
vA0 = 0.0
vB0 = 0.0

t, xA_t, xB_t = movimento_acoplado_2corpos_rk4(xA0, xB0, vA0, vB0)
plt.figure(figsize=(10, 5))
plt.plot(t, xA_t, label='xA(t)', color='blue')
plt.plot(t, xB_t, label='xB(t)', color='orange')
plt.xlabel('Tempo (s)')
plt.ylabel('Posição (m)')
plt.title('Movimento de dois corpos acoplados (caso iii)')
plt.legend()
plt.grid()
plt.show()

#b)
#para o caso i, como u(0) = v(0) (simetricos) as oscilações são idênticas, com a mesma amplitude e fase.
#para o caso ii, como u(0) = -v(0) (inversas) as oscilações são idênticas, mas com a mesma amplitude e fase invertida.
#para o caso iii, como u(0) != v(0) as oscilações são diferentes, com amplitudes e fases diferentes.