import numpy as np
import matplotlib.pyplot as plt

#Exercício 1: Oscilador harmónico forçado

m = 1.0  #kg
k = 1.0  #N/m
b = 0.05 #kg/s
F0 = 7.5 #N
wf = 0.5 #rad/s
x0 = 0.0 #m

#a) Calcule numericamente a lei do movimento, no caso em que a 
# velocidade inicial é nula e a posição inicial 4 m.

v0 = 0.0 #m/s
x0 = 4.0 #m
 
def movimento_oscilador_forcado(m, k, b, F0, wf, x0, v0, t_max=100, n_pts=2000):
    w0 = np.sqrt(k/m)
    beta = b/(2*m)
    wa = np.sqrt(w0**2 - beta**2)

    # Amplitude e fase estacionária
    A = (F0/m) / np.sqrt((wf**2 - w0**2)**2 + (b*wf/m)**2)
    alpha = np.arctan2(b*wf/m, w0**2 - wf**2)

    # Parte transiente
    phi = np.arctan2(-(v0 + x0*beta)/wa, x0)
    Aa = x0 / np.cos(phi)

    # Solução completa
    t = np.linspace(0, t_max, n_pts)
    x_t = (Aa * np.exp(-beta * t) * np.cos(wa * t + phi) +
            A * np.cos(wf * t + alpha))

    return t, x_t

t, x_t = movimento_oscilador_forcado(m, k, b, F0, wf, x0, v0)


plt.figure()
plt.plot(x_t, t)
plt.xlabel('t (s)')
plt.ylabel('x(t) (m)')
plt.title('Lei do movimento')
plt.grid()
plt.show()

#b) Calcule a amplitude do movimento e o seu período no regime estacionário, 
# usando os resultados numéricos.

# ultimos 20% dos dados (regime estacionário)

N = int(0.8 * len(t))
t_stationary = t[N:]
x_stationary = x_t[N:]

# Amplitude
amp = (np.max(x_stationary) - np.min(x_stationary)) / 2
print(f"Amplitude em regime estacionário: {amp:.2f}m")


# Período
imax1 = np.argmax(x_stationary)
tmax1 = t_stationary[imax1]


T_esperado = 2 * np.pi / wf
idx_search = np.where(t_stationary > tmax1 + 0.5 * T_esperado)[0]

if len(idx_search) > 0:
    imax2 = np.argmax(x_stationary[idx_search])
    tmax2 = t_stationary[idx_search][imax2]
    periodo = tmax2 - tmax1
    print(f"Período em regime estacionário: {periodo:.2f}s")
else:
    print("Não foi possivel calcular o periodo")



#c) Repita para valores de 𝜔𝑓 entre 0.2 e 2 rad/s. 
# Faça um gráfico da amplitude em regime estacionário em função de 𝜔𝑓 . 
# Qual a frequência angular 𝜔𝑓 quecorresponde à maior amplitude?

wfs = np.linspace(0.2, 2, 200)
amps = []
periodos = []

for wf in wfs:
    t, x_t = movimento_oscilador_forcado(m, k, b, F0, wf, x0, v0)
    N = int(0.8 * len(t))
    t_stationary = t[N:]
    x_stationary = x_t[N:]

    # Amplitude
    amp = (np.max(x_stationary) - np.min(x_stationary)) / 2
    amps.append(amp)

    #Período: calcula para cada wf
    T_esperado = 2 * np.pi / wf
    imax1 = np.argmax(x_stationary)
    tmax1 = t_stationary[imax1]
    idx_search = np.where(t_stationary > tmax1 + 0.5 * T_esperado)[0]

    if len(idx_search) > 0:
        imax2 = np.argmax(x_stationary[idx_search])
        tmax2 = t_stationary[idx_search][imax2]
        periodo = tmax2 - tmax1
        periodos.append(periodo)
    else:
        periodos.append(np.nan) 

amps = np.array(amps)
periodos = np.array(periodos)

# Gráfico da amplitude em função de wf
plt.figure()
plt.plot(wfs, amps)
plt.xlabel('wf (rad/s)')
plt.ylabel('Amplitude (m)')
plt.title('Amplitude em função de wf')
plt.grid()
plt.show()


max_idx = np.argmax(amps)
print(f"A maior amplitude ocorre para wf = {wfs[max_idx]:.2f} rad/s, com amplitude = {amps[max_idx]:.2}m")