import sys

import os
import torch
import librosa
from open_clip import create_model
from training.data import get_audio_features
from training.data import int16_to_float32, float32_to_int16
from transformers import RobertaTokenizer

tokenize = RobertaTokenizer.from_pretrained("roberta-base")


def tokenizer(text):
    result = tokenize(
        text,
        padding="max_length",
        truncation=True,
        max_length=77,
        return_tensors="pt",
    )
    return {k: v.squeeze(0) for k, v in result.items()}


PRETRAINED_PATH = "/mnt/fast/nobackup/users/hl01486/projects/contrastive_pretraining/CLAP/assets/checkpoints/epoch_top_0_audioset_no_fusion.pt"
WAVE_48k_PATH = "/mnt/fast/nobackup/users/hl01486/projects/contrastive_pretraining/CLAP/assets/audio/machine.wav"


def infer_text():
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    precision = "fp32"
    amodel = "HTSAT-tiny"  # or 'PANN-14'
    tmodel = "roberta"  # the best text encoder in our training
    enable_fusion = False  # False if you do not want to use the fusion model
    fusion_type = "aff_2d"
    pretrained = PRETRAINED_PATH

    model, model_cfg = create_model(
        amodel,
        tmodel,
        pretrained,
        precision=precision,
        device=device,
        enable_fusion=enable_fusion,
        fusion_type=fusion_type,
    )
    # load the text, can be a list (i.e. batch size)
    text_data = ["I love the contrastive learning", "I love the pretrain model"]
    # tokenize for roberta, if you want to tokenize for another text encoder, please refer to data.py#L43-90
    text_data = tokenizer(text_data)

    text_embed = model.get_text_embedding(text_data)
    print(text_embed.size())


def infer_audio():

    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    precision = "fp32"
    amodel = "HTSAT-tiny"  # or 'PANN-14'
    tmodel = "roberta"  # the best text encoder in our training
    enable_fusion = False  # False if you do not want to use the fusion model
    fusion_type = "aff_2d"
    pretrained = PRETRAINED_PATH

    model, model_cfg = create_model(
        amodel,
        tmodel,
        pretrained,
        precision=precision,
        device=device,
        enable_fusion=enable_fusion,
        fusion_type=fusion_type,
    )

    # load the waveform of the shape (T,), should resample to 48000
    audio_waveform, sr = librosa.load(WAVE_48k_PATH, sr=48000)
    # quantize
    audio_waveform = int16_to_float32(float32_to_int16(audio_waveform))
    audio_waveform = torch.from_numpy(audio_waveform).float()
    audio_dict = {}

    # the 'fusion' truncate mode can be changed to 'rand_trunc' if run in unfusion mode
    import ipdb

    ipdb.set_trace()
    audio_dict = get_audio_features(
        audio_dict,
        audio_waveform,
        480000,
        data_truncating="fusion",
        data_filling="repeatpad",
        audio_cfg=model_cfg["audio_cfg"],
    )
    # can send a list to the model, to process many audio tracks in one time (i.e. batch size)
    audio_embed = model.get_audio_embedding([audio_dict])
    print(audio_embed.size())
    import ipdb

    ipdb.set_trace()


if __name__ == "__main__":
    infer_text()
    infer_audio()