import numpy as np
import matplotlib.pyplot as plt
from coef_fourier import abfourier


# Parâmetros físicos
k = 1
k_ = 0.5
m = 1
xA_eq = 1.0
xB_eq = 2.0

# Tempo de simulação
T = 100
dt = 0.01
t = np.arange(0, T, dt)
n_steps = len(t)

# Condições iniciais – caso i
xA0 = xA_eq + 0.3
xB0 = xB_eq + 0.3
vA0 = 0
vB0 = 0

# Inicialização
xA = np.zeros(n_steps)
xB = np.zeros(n_steps)
vA = np.zeros(n_steps)
vB = np.zeros(n_steps)

xA[0] = xA0
xB[0] = xB0
vA[0] = vA0
vB[0] = vB0

# Funções de aceleração
def accel_A(xA, xB):
    return -(k / m) * (xA - xA_eq) - (k_ / m) * (xA - xB)

def accel_B(xA, xB):
    return -(k / m) * (xB - xB_eq) - (k_ / m) * (xB - xA)

# Integração numérica (Euler-Cromer)
for i in range(n_steps - 1):
    aA = accel_A(xA[i], xB[i])
    aB = accel_B(xA[i], xB[i])
    
    vA[i+1] = vA[i] + aA * dt
    xA[i+1] = xA[i] + vA[i+1] * dt
    
    vB[i+1] = vB[i] + aB * dt
    xB[i+1] = xB[i] + vB[i+1] * dt


# Cálculo dos coeficientes de Fourier para xA
it0 = 0
it1 = len(t) - 1
N = 30
omega = 2 * np.pi / T

a_vals = []
b_vals = []
omega_n = []

for n in range(1, N + 1):
    af, bf = abfourier(t, xA, it0, it1, n)
    a_vals.append(af)
    b_vals.append(bf)
    omega_n.append(n * omega)




# Conversão para numpy arrays
omega_vals = np.array(omega_n)
a_vals = np.array(a_vals)
b_vals = np.array(b_vals)

# Plot dos coeficientes a_n e b_n usando linhas e marcadores
plt.figure(figsize=(12, 5))

plt.subplot(1, 2, 1)
plt.plot(omega_vals, a_vals, marker='o', linestyle='-')
plt.title('Coeficientes $a_n$ vs $\\omega_n$')
plt.xlabel('$\\omega_n$ (rad/s)')
plt.ylabel('$a_n$')
plt.grid(True)

plt.subplot(1, 2, 2)
plt.plot(omega_vals, b_vals, marker='s', linestyle='-')
plt.title('Coeficientes $b_n$ vs $\\omega_n$')
plt.xlabel('$\\omega_n$ (rad/s)')
plt.ylabel('$b_n$')
plt.grid(True)

plt.tight_layout()
plt.show()