from speechbrain.pretrained import EncoderClassifier
import torch
import torchaudio
from speechbrain.pretrained.interfaces import foreign_class


class LanguageEmbedding:

    def __init__(self, device=torch.device("cpu")):
        # Use VoxLingua107
        #self.language_id = EncoderClassifier.from_hparams(source="speechbrain/lang-id-voxlingua107-ecapa", savedir="./lang-id-voxlingua107-ecapa")
        self.language_id = foreign_class(source="TalTechNLP/voxlingua107-xls-r-300m-wav2vec", pymodule_file="encoder_wav2vec_classifier.py", classname="EncoderWav2vecClassifier", hparams_file='inference_wav2vec.yaml', savedir="tmp")

        # Use pretrained model:
        #print("using standard pretrained model in ./Preprocessing/Language_embedding.py")
        # self.language_id = EncoderClassifier.from_hparams(source="speechbrain/lang-id-commonlanguage_ecapa",
        #                                                   run_opts={
        #                                                       "device": str(device)},
        #                                                   savedir="pretrained_models/lang-id-commonlanguage_ecapa")

        #  Use 45+WASS model:
        # print("attention this is only the pretrained model!!!")
        # self.language_id = EncoderClassifier.from_hparams(source="/data/vokquant/speechbrain/recipes/CommonLanguage/lang_id/results/ECAPA-TDNN/1986/save/ckpt/",
        #                                                   run_opts={"device": str(device)},
        #                                                   savedir="/data/vokquant/IMS-Toucan_lang_emb/Preprocessing/pretrained_models/lang-id-commonlanguage_ecapa")
        
    def get_language_embedding(self, input_waves=None):
        #print("attention this is only the pretrained model!!!")
        # print(input_waves.size())
        # embeddings = self.language_id.encode_batch(input_waves)
        # wave,sr = soundfile.read(wav_file)
        # wave = torch.from_numpy(wave)
        embeddings = self.language_id.encode_batch(wavs=input_waves)
        # print text_lab for 1 file of batch
        # out_prob, score, index, text_lab = self.language_id.classify_batch(input_waves)
        #print(text_lab)
        #print((embeddings.shape()))
        return embeddings

    def get_emb_from_path(self, path_to_wavfile=None):
        #print("attention this is only the pretrained model!!!")
        audio = self.language_id.load_audio(path_to_wavfile)
        out_prob, score, index, text_lab = self.language_id.classify_file(path_to_wavfile)
        print(text_lab)
        #print(self.language_id.encode_batch(audio))
        print(self.language_id.encode_batch(audio).shape)
        # get number of dimensions of size
        #print(self.language_id.encode_batch(audio).squeeze(0))
        #return self.language_id.encode_batch(audio)
        #print(len(self.language_id.encode_batch(audio).shape))
        if len(self.language_id.encode_batch(audio).shape) == 3:
            return self.language_id.encode_batch(audio).squeeze(0)
        else:
            return self.language_id.encode_batch(audio)
        
if __name__ == '__main__':
    #print('hi')
    # load example wav file
    path_to_wavfile = "/nas/projects/vokquant/data/aridialect/aridialect_wav16000/hpo_at_berlin_004.wav"
    #path_to_wavfile = "/nas/projects/vokquant/data/aridialect/aridialect_wav16000/hoi_goi_goi_004.wav"
    # path_to_wavfile = "speechbrain/lang-id-commonlanguage_ecapa/example-fr.wav"
    # get emb from path
    lang_emb = LanguageEmbedding().get_emb_from_path(path_to_wavfile)