import argparse
import torch
import librosa
import numpy as np
from torch import autocast
from contextlib import nullcontext
import os

from models.mn.model import get_model as get_mobilenet
from models.dymn.model import get_model as get_dymn
from models.ensemble import get_ensemble_model
from models.preprocess import AugmentMelSTFT
from helpers.utils import NAME_TO_WIDTH, labels


def audio_tagging(args):
    """
    Running Inference on an audio clip.
    """
    model_name = args.model_name
    pre_trained_model_type = args.pre_trained_model_type
    pre_trained_model_path = args.pre_trained_model_path
    device = torch.device('cuda') if args.cuda and torch.cuda.is_available() else torch.device('cpu')
    audio_path = args.audio_path
    sample_rate = args.sample_rate
    window_size = args.window_size
    hop_size = args.hop_size
    n_mels = args.n_mels
    
    if os.path.isfile(pre_trained_model_path):
        if pre_trained_model_type.startswith("dymn"):
            model = get_dymn(width_mult=NAME_TO_WIDTH(pre_trained_model_type), pretrained_name=pre_trained_model_path,
                                strides=args.strides)
        else:
            model = get_mobilenet(width_mult=NAME_TO_WIDTH(pre_trained_model_type), pretrained_name=pre_trained_model_path,
                                strides=args.strides, head_type=args.head_type)
    else:
        if model_name.startswith("dymn"):
            model = get_dymn(width_mult=NAME_TO_WIDTH(model_name), pretrained_name=model_name,
                                    strides=args.strides)
        else:
            model = get_mobilenet(width_mult=NAME_TO_WIDTH(model_name), pretrained_name=model_name,
                                    strides=args.strides, head_type=args.head_type)
    model.to(device)
    model.eval()

    # model to preprocess waveform into mel spectrograms
    mel = AugmentMelSTFT(n_mels=n_mels, sr=sample_rate, win_length=window_size, hopsize=hop_size)
    mel.to(device)
    mel.eval()

    (waveform, _) = librosa.core.load(audio_path, sr=sample_rate, mono=True)
    waveform = torch.from_numpy(waveform[None, :]).to(device)

    # our models are trained in half precision mode (torch.float16)
    # run on cuda with torch.float16 to get the best performance
    # running on cpu with torch.float32 gives similar performance, using torch.bfloat16 is worse
    with torch.no_grad(), autocast(device_type=device.type) if args.cuda else nullcontext():
        spec = mel(waveform)
        preds, _ = model(spec.unsqueeze(0))
    
    # Convert logits to predicted class
    preds = preds.float().squeeze()  # Ensure correct type and shape
    predicted_class = torch.argmax(preds).item()

    # Print detected class
    print("************* Detected Acoustic Event: *****************")
    print(f"Predicted Class: {labels[predicted_class]} (Class ID: {predicted_class})")
    print("*******************************************************")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Example of parser. ')
    # model name decides, which pre-trained model is loaded
    parser.add_argument('--pre_trained_model_type', type=str, default='dymn04')
    parser.add_argument('--pre_trained_model_path', type=str, default='')
    parser.add_argument('--model_name', type=str, default='mn10_as')
    parser.add_argument('--strides', nargs=4, default=[2, 2, 2, 2], type=int)
    parser.add_argument('--head_type', type=str, default="mlp")
    parser.add_argument('--cuda', action='store_true', default=False)
    parser.add_argument('--audio_path', type=str, required=True)

    # preprocessing
    parser.add_argument('--sample_rate', type=int, default=32000)
    parser.add_argument('--window_size', type=int, default=800)
    parser.add_argument('--hop_size', type=int, default=320)
    parser.add_argument('--n_mels', type=int, default=128)


    args = parser.parse_args()

    audio_tagging(args)
