import os
import numpy as np
import librosa
import soundfile as sf
from PIL import Image

def ensure_dir(path):
    if not os.path.exists(path):
        os.makedirs(path)

def convert_wav_to_spectrogram(audio_path, output_image_path):
    """
    Charge l'audio, génère un spectrogramme Mel, normalise l'image, 
    la redimensionne en 224x224 et la sauvegarde (sans axes).
    """
    # 1. Charger l'audio
    y, sr = librosa.load(audio_path, sr=None)
    
    # 2. Spectrogramme de Mel
    melspec = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128)
    
    # 3. Échelle Logarithmique (Décibels)
    S_db = librosa.power_to_db(melspec, ref=np.max)
    
    # 4. Normalisation entre 0 et 255 pour image
    S_db_min = S_db.min()
    S_db_max = S_db.max()
    if S_db_max == S_db_min:
        img_array = np.zeros_like(S_db, dtype=np.uint8)
    else:
        norm_spec = (S_db - S_db_min) / (S_db_max - S_db_min)
        img_array = (norm_spec * 255).astype(np.uint8)
    
    # On inverse l'axe Y pour avoir les basses fréquences en bas de l'image
    img_array = np.flip(img_array, axis=0)
    
    # 5. Création de l'image, conversion RGB et Resize stricte
    img = Image.fromarray(img_array)
    # Convertir en RGB pour garder la compatibilité ResNet
    img = img.convert('RGB')
    img = img.resize((224, 224), Image.Resampling.BILINEAR)
    
    # 6. Sauvegarde
    img.save(output_image_path, "JPEG")
    print(f"Spectrogramme sauvegardé : {output_image_path}")

def create_test_audio(output_path):
    """Génère un fichier .wav de test basique (onde sinusoïdale)."""
    sr = 22050
    duration = 2.0 # 2 secondes
    t = np.linspace(0, duration, int(sr * duration), endpoint=False)
    # Sine wave à 440 Hz (Note La)
    audio_data = 0.5 * np.sin(2 * np.pi * 440.0 * t)
    
    sf.write(output_path, audio_data, sr)
    print(f"Fichier audio de test généré : {output_path}")

def main():
    # Chemins
    raw_dir = "data/audio/raw"
    spec_dir = "data/audio/spectrograms"
    
    # Création des dossiers
    ensure_dir(raw_dir)
    ensure_dir(spec_dir)
    print("Dossiers créés avec succès.")
    
    test_audio_path = os.path.join(raw_dir, "test.wav")
    test_spec_path = os.path.join(spec_dir, "test.jpg")
    
    # Génération
    create_test_audio(test_audio_path)
    
    # Conversion
    convert_wav_to_spectrogram(test_audio_path, test_spec_path)
    
    if os.path.exists(test_spec_path):
         print("\nSUCCĖS : L'image du spectrogramme a bien été générée sur le disque !")
    else:
         print("\nERREUR : Le spectrogramme n'a pas pu être sauvegardé.")

if __name__ == "__main__":
    main()