import numpy as np
import scipy.fftpack as FFT
from matplotlib import pyplot as plt

f0=1
NbPeriodes=5
duree=NbPeriodes/f0

def mono1(t):
    return np.cos(2*np.pi*t)
    
def mono6(t):
    return np.cos(2*np.pi*6*t)
    
def polychromatique(t):
    freq=[1,2,4,6]
    amplitude=[1,2/3,8/27,32/243]
    phase=[0,0,0,0]
    som=0
    for i in range(len(freq)):
        som+=amplitude[i]*np.cos(freq[i]*2*np.pi*t+phase[i])
    return som

def creneau(t):
    t1=t%1
    if t1<1/4 or t1>3/4:
        return -1
    return 1
    
def triangle(t):
    t1=t%1
    if t1<1/2:
        return 1-4*t1
    return 4*t1-3
    
def echantillonnage(s,fe):
    N=int(duree*fe)
    dt=duree/N
    temps=[i*dt for i in range(N)]
    echantillon=[s(t) for t in temps]
    return temps,echantillon
    
def spectre(s,fe,reel=False):
    temps,echantillon=echantillonnage(s,fe)
    N=len(echantillon)
    if reel:
        TFD=np.abs(np.fft.rfft(echantillon))
        freq=np.linspace(0,fe/2,N//2+1)
    else:
        TFD=np.abs(FFT.fft(echantillon))
        freq=np.linspace(0,fe,N)
    fonction="Fonction "+s.__name__
    plt.figure(fonction)
    plt.subplot(121) 
    T=np.linspace(0,duree,500)
    plt.plot(T,[s(t) for t in T],color="k")
    plt.plot(temps,echantillon,"or") 
    plt.xlabel("temps (s)")   
    plt.subplot(122)
    plt.bar(freq,TFD,width=0.1)
    plt.xlabel("fréquences (Hz)")  
    plt.show()
    
#spectre(mono1,10)
#spectre(mono1,10,True)

#spectre(mono6,10)
#spectre(mono6,10,True)

#spectre(polychromatique,15)
#spectre(polychromatique,11)

#spectre(polychromatique,15,True)
#spectre(polychromatique,11,True)

#spectre(creneau,15)
#spectre(creneau,15,True)