import torch
import torch.nn as nn
from tqdm import tqdm
from lm_eval import evaluator
from datasets import load_dataset
from transformers import pipeline
from evaluate import load as load_metric
from lm_eval.tasks import initialize_tasks
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.models.whisper.english_normalizer import BasicTextNormalizer


def get_device():
    if torch.backends.mps.is_available():
        return "mps"
    elif torch.cuda.is_available():
        return "cuda:0"
    else:
        return "cpu"


def evaluate_perplexity(model, tokenizer):
    def _perplexity(nlls, n_samples, seqlen):
        return torch.exp(torch.stack(nlls).sum() / (n_samples * seqlen))

    # load and prepare dataset
    data = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
    data = tokenizer("\n\n".join(data["text"]), return_tensors="pt")
    data = data.input_ids.to(model.device)

    seqlen = 2048
    model = model.eval()
    n_samples = data.numel() // seqlen

    nlls = []

    with tqdm(range(n_samples), desc="Perplexity -") as progress_bar:
        for i in progress_bar:
            start_index = i * seqlen
            end_index = (i + 1) * seqlen
            batch = data[:, start_index:end_index].to(model.device)
            with torch.no_grad():
                logits = model(batch).logits
            shift_logits = logits[:, :-1, :].contiguous().float()
            shift_labels = data[:, start_index:end_index][:, 1:]
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(
                shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)
            )
            neg_log_likelihood = loss.float() * seqlen
            nlls.append(neg_log_likelihood)

            curr_ppl = _perplexity(nlls, i + 1, seqlen)
            progress_bar.set_description(f"Perplexity {curr_ppl:.3f}")

    ppl = _perplexity(nlls, n_samples, seqlen)

    return ppl.item()


def eval_librispeech(model_id, num_samples=100, batch_size=4):
    try:
        import jiwer, librosa, soundfile
    except ImportError:
        print("Please install the following: pip install jiwer librosa soundfile")

    dataset = load_dataset("librispeech_asr", "clean", split="test", streaming=True)

    # Load the Whisper model pipeline for automatic speech recognition
    pipe = pipeline(
        task="automatic-speech-recognition",
        model=model_id,
        batch_size=batch_size,
        device=get_device(),
        torch_dtype=torch.float16,
    )

    # Word normalizer
    normalizer = BasicTextNormalizer()

    # Load the WER metric
    wer_metric = load_metric("wer")

    texts = []
    audio = []
    for i, data in tqdm(enumerate(dataset), total=num_samples, desc="Loading dataset"):
        if len(audio) == num_samples:
            break
        audio.append(data["audio"])
        texts.append(data["text"])

    references = []
    predictions = []

    with tqdm(range(0, num_samples, batch_size), desc="Word Error Rate: -") as pbar:
        for i in pbar:
            batch_audio = audio[i : i + batch_size]
            batch_texts = texts[i : i + batch_size]

            # inference
            results = pipe(batch_audio, batch_size=len(batch_audio))

            # normalize text
            normalized_predictions = [normalizer(result["text"]) for result in results]
            normalized_texts = [normalizer(text) for text in batch_texts]

            predictions.extend(normalized_predictions)
            references.extend(normalized_texts)

            # word error rate computation
            wer = (
                wer_metric.compute(predictions=predictions, references=references) * 100
            )
            pbar.set_description(f"Word Error Rate: {wer:.3f}%")


def eval_mmlu(
    model_path="gpt2",
    num_fewshot=1,
    batch_size=1,
    device="cuda:0",
    task_use_pretrained=False,
):
    try:
        import vllm

        VLLM_INSTALLED = True
    except ImportError:
        VLLM_INSTALLED = False

    initialize_tasks(verbosity="DEBUG")

    if VLLM_INSTALLED:
        model = "vllm"
        model_args = dict(
            pretrained=model_path,
            max_model_len=2048,
            dtype="float16",
            trust_remote_code=True,
        )

        if not task_use_pretrained:
            model_args["quantization"] = "awq"
    else:
        model = "hf"
        model_args = dict(
            pretrained=model_path,
            device_map_option=device,
            dtype="float16",
            trust_remote_code=True,
        )
    model_args = ",".join([f"{k}={v}" for k, v in model_args.items()])

    results = evaluator.simple_evaluate(
        model=model,
        model_args=model_args,
        tasks=["mmlu"],
        num_fewshot=num_fewshot,
        batch_size=batch_size,
        device=device,
        log_samples=False,
    )

    print(evaluator.make_table(results))


if __name__ == "__main__":
    ### PERPLEXITY
    # model_path = 'mistralai/Mistral-7B-Instruct-v0.1'
    # model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
    # tokenizer = AutoTokenizer.from_pretrained(model_path)
    # evaluate_perplexity(model, tokenizer)

    ### WORD ERROR RATE
    # model_id = "distil-whisper/distil-small.en" # 3.594
    model_id = "distil-whisper/distil-medium.en"  # 3.436
    eval_librispeech(model_id)