import numpy as np
import matplotlib.pyplot as plt

#Exercício 2: Oscilador Acoplado: Solução Teórica

#O movimento do oscilador acoplado do exercício anterior pode ser descrito como
#uma sobreposição de dois modos normais, cada um sendo um movimento
#harmónico simples

m = 1 #kg
kl = 0.5  #N/m
k= 1  #N/m
xA_eq = 1.0 #m
xB_eq = 2.0 #m
ω1 = np.sqrt(k/m)
ω2 = np.sqrt((k+2*kl)/m)


tf = 100
dt = 0.01
n_passos = int(tf / dt)
t  = np.linspace(0, tf, n_passos)



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):

    dt = 0.1
    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]



def compara_analitica_numerica(xA0, xB0, nome, xA_eq, xB_eq,
                                ω1, ω2, func,t, t_max):
    #coeficientes
    u0 = xA0 - xA_eq
    v0 = xB0 - xB_eq
    C1 = 0.5 * (u0 + v0)
    C2 = 0.5 * (u0 - v0)
    A1 = abs(C1)
    A2 = abs(C2)
    phi1 = 0.0 if C1 >= 0 else np.pi
    phi2 = 0.0 if C2 >= 0 else np.pi


    xA_teorica = xA_eq + A1*np.cos(ω1*t + phi1) + A2*np.cos(ω2*t + phi2)
    xB_teorica = xB_eq + A1*np.cos(ω1*t + phi1) - A2*np.cos(ω2*t + phi2)

    
    t_num, xA_num, xB_num = func(xA0, xB0, vA0, vB0, t_max)

    plt.figure(figsize=(8,3))
    plt.plot(t, xA_teorica, 'k--', label='A teorica')
    plt.plot(t, xB_teorica, 'r--', label='B teorica')
    plt.plot(t_num, xA_num,'k',   label='A numerica')
    plt.plot(t_num, xB_num,'r',   label='B numerica')
    plt.title(f'Caso {nome}: A1={A1:.2f}, A2={A2:.2f}')
    plt.xlabel('t (s)'); plt.ylabel('x (m)')
    plt.legend(); plt.tight_layout()
    plt.show()

    print(f'Caso {nome}: A1={A1:.2f}, A2={A2:.2f}, phi1={phi1:.2f}, phi2={phi2:.2f}')


#i) 
xA0 = xA_eq + 0.3
xB0 = xB_eq + 0.3
vA0 = 0.0
vB0 = 0.0
compara_analitica_numerica(xA0, xB0, 'i', xA_eq, xB_eq,
                            ω1, ω2, movimento_acoplado_2corpos_rk4,
                            t, tf)

#ii)
xA0 = xA_eq + 0.3
xB0 = xB_eq - 0.3
vA0 = 0.0
vB0 = 0.0
compara_analitica_numerica(xA0, xB0, 'ii', xA_eq, xB_eq,
                            ω1, ω2, movimento_acoplado_2corpos_rk4,
                            t, tf)

#iii)
xA0 = xA_eq + 0.3
xB0 = xB_eq - 0.1
vA0 = 0.0
vB0 = 0.0
compara_analitica_numerica(xA0, xB0, 'iii', xA_eq, xB_eq,
                            ω1, ω2, movimento_acoplado_2corpos_rk4,
                            t, tf)