#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu Feb 27 09:37:23 2025

@author: home
"""

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

c = 343
f_e = 40e3
A0 = 0.76

def a_th(X0 = 3/100,d0 = 49/1000):
    
    ''' paramètres initiaux '''
    m = 1/5
    k_r = 90
    N = 1000000
    
    t = np.linspace(0,1,N)
    
    w = 2*np.pi*f_e #pulsation capteur 
    k = w/c
    
    ''' etude du mouvement du ressort '''
    w0 = np.sqrt(k_r/m) #pulsation
    X_r = -X0*np.cos(w0*t) #position
    V_r = w0*X0*np.sin(w0*t) #vitesse
    
    ''' frequence, phase, amplitude '''
    f_ra = f_e/(1-np.abs(V_r)/c) #approche
    f_re = f_e/(1+np.abs(V_r)/c) #eloignement
    
    dist_cc = d0 + X_r
    
    def A(x):
        return A0/x*(1+13.51/100*np.cos(k*x/2))
    
    S1 = np.where(0 <= V_r, 
                  A(dist_cc) * np.cos(2 * np.pi * f_re * t-k*dist_cc), 
                  A(dist_cc) * np.cos(2 * np.pi * f_ra * t-k*dist_cc)) #signal modifié par le ressort

    fig1, ax = plt.subplots(figsize=(10,4))
    ax.plot(t, S1, 'pink')
    #ax.plot(t, Env, 'r', lw=1.2, label='Enveloppe des battements')
    ax.set(title='Signal au repos', xlabel="temps (s)", ylabel="Amplitude",xlim=(0,0.25))  
    #ax.legend(loc='upper right', fontsize=10.5)
    
    with open("donnees theoriques.txt", "w", encoding="utf-8") as fichier:
        fichier.write("temps;Somme des signaux\n")   
        for temps, signal in zip(t, S1):
            fichier.write(f"{temps};{signal}\n")
    
def a_fft(x=str): 
    
    t, data = np.loadtxt(x,delimiter=';',skiprows=1,unpack=True)
    dt = t[1]-t[0]
    print(f"Nombre d'échantillons : {len(data)}\nDurée du signal : {t[-1]} s")

    f_nu = spfft.rfft(data) # calcul de la transformée de fourier
    freq = spfft.rfftfreq(len(data), dt) # calcul des fréquences de Fourier
    
    fig, ax = plt.subplots(figsize=(10,4))   
    ax.plot(freq, np.abs(f_nu))
    ax.set(xlim=(39000,41000),title='FFT du signal',xlabel="fréquence (Hz)",ylabel='Amplitude')
    
def a_stft(x=str):
    
    t, data = np.loadtxt(x,delimiter=';',skiprows=1,unpack=True)
    rate = 1/(t[1]-t[0])
    print(f"Nombre d'échantillons : {len(data)}\nDurée du signal : {t[-1]} s")
    
    N = 100 
    s_moy = np.sqrt(2) * np.convolve(np.abs(data), np.ones(N)/N, mode='same')
    
    f, t_s, Zxx = stft(data/s_moy, rate, nperseg=16384) 
        
    freq=[]   
    for i in range(Zxx.shape[1]):
        I = np.argmax(np.abs(Zxx[:,i]),axis=0)
        F = f[I]
        freq.append(F)
            
    fig, ax = plt.subplots(figsize=(10, 4))
    img = ax.pcolormesh(t_s,f,np.abs(Zxx), cmap='coolwarm')
    ax.set(title="Spectrogramme du signal sonore", ylabel="Fréquence (Hz)", xlabel="Temps (s)",ylim=(38000,42000))
    fig.colorbar(img,label="Amplitude")
    ax.scatter(t_s,freq,color='yellow',s=3)
    
    def vitesse_exp(f):
        f = np.array(f)
        return c*np.abs(f_e-f)/f_e
    
    V = vitesse_exp(freq)
    fig1,ax=plt.subplots(figsize=(10,4))
    ax.plot(t_s,V,color='mediumslateblue',label='Vitesse expérimentale')
    ax.set(title="Vitesses", ylabel=r"Vitesse ($\mathrm{m \cdot s^{-1}}$)", xlabel="Temps (s)")
    #ax.legend(loc='upper right', fontsize=10.5)