import torch
import librosa
import time
from transformers import set_seed, Wav2Vec2ForCTC, AutoProcessor
import numpy as np

device = 0 if torch.cuda.is_available() else "cpu"

def goai_stt(fichier):
    """
    Transcrire un fichier audio donné.
    
    Paramètres
    ---------- 
    fichier: str | tuple[int, np.ndarray]
        Le chemin d'accès au fichier audio ou le tuple contenant le taux d'échantillonnage et les données audio.
        
    Return
    ---------- 
    transcript: str
        Le texte transcrit.
    """

    print("Fichier entré en entréé ---------> ", fichier)
    
    if fichier is None:
        raise ValueError("Le fichier audio est manquant.")

    ### assurer reproducibilité
    set_seed(2024)
    
    start_time = time.time()
    
    ### charger le modèle de transcription 
    model_id = "anyantudre/wav2vec2-large-mms-1b-mos-V1"

    processor = AutoProcessor.from_pretrained(model_id)
    model = Wav2Vec2ForCTC.from_pretrained(model_id, target_lang="mos", ignore_mismatched_sizes=True).to(device)

    if isinstance(fichier, str):
        ### preprocessing de l'audio à partir d'un fichier
        signal, sampling_rate = librosa.load(fichier, sr=16000)
    else:
        ### preprocessing de l'audio à partir d'un tableau numpy
        sampling_rate, signal = fichier

    # Convert the signal to float32
    signal = signal.astype(np.float32)

    inputs = processor(signal, sampling_rate=16000, return_tensors="pt", padding=True).to(device)
    
    ### faire l'inference
    with torch.no_grad():
        outputs = model(**inputs).logits
    
    pred_ids = torch.argmax(outputs, dim=-1)[0]
    transcription = processor.decode(pred_ids)

    print("Temps écoulé: ", int(time.time() - start_time), " secondes")
    return transcription