#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Sep 21 17:25:37 2023

@author: eddiesaudrais
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy import signal
from numpy.fft import fft

#------------------------------------------------------------------------------

# signaux de base. 

f = 1       # fréquence du signal (à modifier)

def sinus(t): 
    return np.sin(2. * np.pi * f * t)

def triangle(t):
    return signal.sawtooth(2. * np.pi * f * t,0.5)

def dent_scie(t):
    return signal.sawtooth(2. * np.pi * f * t,1)

def carre(t):
    return signal.square(2. * np.pi * f * t,.5)

#------------------------------------------------------------------------------

# S          fonction définissant le signal
# F_ech      fréquence d'échantillonnage
# duree      durée totale de l'échantillon
# (F_max)    (optionnel) fréquence maximale de tracé du spectre
f
def spectre(S,F_ech,duree,*args):
    t = np.arange(0,duree,1/F_ech)                      # on construit le tableau des instants d'échantillonnage
    ech = S(t)                                          # on construit l'échantillon
    tfd = fft(ech)                                      # on calcule le spectre par FFT
    N = len(tfd)                                        # on construit le tableau des fréquences 
    freq = np.zeros(N)
    for k in range(N):
        freq[k] = 1./duree*k
    spectre = np.absolute(tfd)*2./N                     # norme des composantes du spectre, en normalisant
    spectre[0] = .5*spectre[0]                          # normalisation de la composante continue
    if len(args)>0:                                     # argument optionnel : on trace le spectre
        f_max = args[0]                                 # jusqu'à la fréquence f_max
    else:
        f_max = F_ech
    N_max = int(N*f_max/F_ech)
    plt.close()
    plt.stem(freq[:N_max],spectre[:N_max],markerfmt=" ",basefmt=" ")    # tracé du spectre
    plt.xlabel("frequence")
    plt.ylabel("amplitude")
    plt.axhline(0)
    plt.title("spectre calculé par FFT avec $F_\mathrm{éch}=$"+str(F_ech)+" Hz")
    plt.show()

    
#------------------------------------------------------------------------------

# fentre sélectionnée en 4e argument

def spectre_fenetre(signal,F_ech,duree,fenetre,*args):
    t = np.arange(0,duree,1/F_ech)      # on construit le tableau des instants d'échantillonnage
    ech = signal(t)                     # on construit l'échantillon
    M = len(ech)                        # nombre de points de l'échantillon
    if fenetre == 1:
        window = np.hamming(M)          # fenêtre de Hamming
        nom = "avec fenêtre de Hamming"
    elif fenetre == 2:
        window = np.hanning(M)          # fenêtre de Hanning 
        nom = "avec fenêtre de Hanning"
    elif fenetre == 3:
        window = np.blackman(M)         # fenêtre de Blackman 
        nom = "avec fenêtre de Blackman"
    elif fenetre == 4:
        window = np.bartlett(M)         # fenêtre de Bartlett
        nom = "avec fenêtre de Bartlett"
    elif fenetre == 5:
        window = np.kaiser(M,14)        # fenêtre de Kaiser (coeff. modifiable)
        nom = "avec fenêtre de Kaiser"
    else:
        window = 1.
        nom = "sans fenêtre"
    ech_fenetre = window*ech
    tfd = fft(ech_fenetre)              # on calcule le spectre par FFT
    N = len(tfd)                        # on construit le tableau des fréquences 
    freq = np.zeros(N)
    for k in range(N):
        freq[k] = 1./duree*k
    spectre = np.absolute(tfd)*2./N     # norme des composantes du spectre, en normalisant
    spectre[0] = .5*spectre[0]          # normalisation de la composante continue
    if len(args)>0:                     # argument optionnel : on trace le spectre
        f_max = args[0]                 # jusqu'à la fréquence f_max
    else:
        f_max = F_ech
    N_max = int(N*f_max/F_ech)
    plt.close()
    plt.stem(freq[:N_max],spectre[:N_max],markerfmt=" ",basefmt=" ")    # tracé du spectre
    plt.xlabel("fréquence")
    plt.ylabel("amplitude")
    plt.title("spectre calculé par FFT "+str(nom))
    plt.show()