import numpy as np
import matplotlib.pyplot as plt

#Exercício 2: Oscilador Não Harmónico Caótico, com Metodo de Runge-Kutta


m = 1.0  #kg
alpha = 1  #N/m^3
k = 0.2 #N/m
b = 0.01 #kg/s
F0 = 5 #N
wf = 0.6 #rad/s


#a) ) Use o método de Runge-Kutta da 4ª ordem para calcular numericamente a 
# lei do movimento,de 0 até 50s

v0 = 0.0 #m/s
x0 = 1 #m

#f = ma
def acelera(t, x, vx):
    return (-k*x - 4*alpha*x**3 - b*vx + F0*np.cos(wf*t)) / m


def rk4_x_vx(t, x, vx, acelera, dt):
    ax1 = acelera(t, x, vx)
    c1v = ax1*dt
    c1x = vx*dt

    ax2 = acelera(t + dt/2., x + c1x/2., vx + c1v/2.)
    c2v = ax2*dt
    c2x = (vx + c1v/2.)*dt

    ax3 = acelera(t + dt/2., x + c2x/2., vx + c2v/2.)
    c3v = ax3*dt
    c3x = (vx + c2v/2.)*dt

    ax4 = acelera(t + dt, x + c3x, vx + c3v)
    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



dt = 0.01 
t0 = 0
tf = 50

n = int((tf-t0)/dt)
t_vec = np.linspace(t0, tf, n+1)


x = np.zeros(n+1)
vx = np.zeros(n+1)
x[0] = x0
vx[0] = v0


for i in range(n):
    x[i+1], vx[i+1] = rk4_x_vx(t_vec[i], x[i], vx[i], acelera, dt)

# Gráfico da lei do movimento
plt.plot(t_vec, x)
plt.xlabel('t (s)')
plt.ylabel('x(t) (m)')
plt.title('Oscilador Nao Harmonico Quartico')
plt.grid()
plt.show()

#b) dt > 0.01 

#c) 
# tf = 100
# x0 = 1.0001

#a fase final e muito diferente da fase final do a), que indica o comportamente
#caotico do sistema (muito dependente das condicoes iniciais)

#d)
# Grafico do espaço de fase (trajetoria aleatoria)
plt.figure(figsize=(6, 6))
plt.plot(x, vx)
plt.xlabel('x (m)')
plt.ylabel('vₓ (m/s)')
plt.title('Espaço de Fase: vx(t) - x(t)')
plt.grid()
plt.show()