import numpy as np
import scipy.integrate as integr
from matplotlib import pyplot as plt

E0=1
f0=100
T0=1/f0
omega0=2*np.pi*f0

def monochromatique(t):
    return E0*np.cos(omega0*t)
    
def polychromatique(t):
    freq=[1,5]
    amplitude=[1,1/3]
    phase=[0,0]
    som=0
    for i in range(len(freq)):
        som+=E0*amplitude[i]*np.cos(freq[i]*omega0*t+phase[i])
    return som

def creneau(t):
    t1=t%T0
    if t1<T0/4 or t1>3*T0/4:
        return -E0
    return E0
    
def triangle(t):
    t1=t%T0
    if t1<T0/2:
        return -4*E0/T0*(t1-T0/4)
    return 4*E0/T0*(t1-3*T0/4)

def coeffFourier(s,n):
    if n==0:
        return integr.quad(s,0,T0)[0]/T0,0
    else:
        #calcul à compléter
        return an,bn

def energie(s):
    def sCarre(t):
        return s(t)*s(t)
    return integr.quad(sCarre,0,T0)[0]/T0
        
def serieFourier(coeff,T):
    som=0
    for i in range(len(coeff)):
        som+=coeff[i][0]*np.cos(i*omega0*T)
        som+=coeff[i][1]*np.sin(i*omega0*T)
    return som
    
def spectre(s,N):
    coeff=[coeffFourier(s,n) for n in range(N+1)]
    fonction="Fonction "+s.__name__
    plt.figure(fonction)
    plt.subplot(121)
    T=np.linspace(-T0/6,7*T0/6,500)
    plt.plot(T,[s(t) for t in T],color="k")
    SF=serieFourier(coeff,T)
    plt.plot(T,SF,color="r")
    plt.subplot(122)
    modules=[np.abs(c[0]+1j*c[1]) for c in coeff]
    plt.bar(list(range(N+1)),modules)
    som=modules[0]**2
    for m in modules[1:]:
        som+=m*m/2
    print(som,energie(s))
    plt.show()