import numpy as np
import matplotlib.pyplot as plt
plt.ion() # esto habilita el interactive plotting, más fácil para trabajar en consola
from numpy import pi as pi

# Definimos los parámetros del problema
t_c = 1
g = 1

def sol_posta(t):
    """
    Función que devuelve la solución exacta para y(t)
    """
    v_0 = g*t_c/2

    return v_0*t - g/2*t**2 

def componente(t,n):
    """
    Función que devuelve la componente de Fourier de y(t)
    en función de n, para cada t.
    """
   
    a_n = 4*t_c**2/pi**3 * g/n**3
    
    return a_n*np.sin(n*pi*t/t_c)

def fourier(t, N):
    """
    Función que devuelve la suma parcial de Fourier de y(t)
    para un dado N, para cada t.
    """
    
    suma = np.zeros(len(t))
 
    # hago la suma sobre los n impares, usando el tercer argumento de range()
    for n in range(1,N+2,2):

        suma += componente(t,n)

    return suma

# Genero un array de tiempos
tiempos = np.linspace(0,t_c,1000)

# Primero grafico la solución exacta
plt.figure()
plt.plot(tiempos, sol_posta(tiempos),'k', label = 'solución exacta', lw = 3)

# Defino aquí el máximo orden del desarollo de Fourier
N_max = 9

for N in range(1,N_max+2,2):

    plt.plot(tiempos, fourier(tiempos,N),'--', label = '$N={}$'.format(N))

plt.xlabel('$t$', fontsize = 15)
plt.ylabel('$y(t)$', fontsize = 15)
plt.grid(True)
plt.legend()
plt.title('Aproximación a la trayectoria por sumas parciales', fontsize = 15)

# Por último, guarda la figura como imagen
plt.savefig('problema_3_5.png')