import os
from torcheval.metrics import WordErrorRate
import json
import argparse
import wandb
import whisper
import tqdm
from suno_utils.audio import Audio
import re
from suno_utils.utils.lyrics import remove_speakers
from suno_utils.utils.text import normalize_whitespace


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str)
    args = parser.parse_args()
    return args


def clean_text(text):
    """General text cleaning."""
    text = "\n" + text
    text = text.replace("’", "'").lower()
    text = text.replace('"', "").lower()
    text = re.sub(r"\[.+?\]", " ", text)  # tags
    text = re.sub(r"\n.+?\:", " ", text)  # new line ends with :
    text = re.sub(r"\n.+?\：", " ", text)  # new line ends with :
    text = re.sub(r"\n\(.+?\)", " ", text)  # new line with ()
    text = re.sub(r"[\d]", " ", text)  # digits
    text = re.sub(r"▁", "", text)  # special stuff
    text = remove_speakers(text)
    text = re.sub(r"[^\w\'\s]", " ", text)  # keep only the words
    text = normalize_whitespace(text)
    return text


if __name__ == "__main__":
    avialbe_device = f"cuda:{os.environ['CUDA_VISIBLE_DEVICES']}"
    print(avialbe_device)
    input_args = parse_args()
    model = whisper.load_model("large-v2")  # large-v2 takes like 20 sec per song
    model.to("cuda")
    print("finish loading whisper")
    wandb.init(project="hoot_evals", name=f"whisper_large_v2")
    validation_datasets = {
        # "english": "/home/tony/Data/Hoot/en_test_manifest_norm.json",
        # "multi": "/home/tony/Data/Hoot/all_test_manifest.json",
        # "multi_balanced": "/home/tony/Data/Hoot/multi_balanced_test_manifest.json",
        "v3_balanced": "/home/tony/Data/Hoot/v3_validation_set.json",
    }
    for validation_set, validation_path in validation_datasets.items():
        # if validation_set != input_args.dataset:
        #     continue
        whisper_results = {}
        log_dict = {}
        val_paths = []
        val_truth_text = []
        with open(validation_path, "r") as fp:
            for l in fp:
                val_meta = json.loads(l)
                val_paths.append(val_meta["audio_filepath"])
                val_truth_text.append(val_meta["text"])
        whisper_outs = []
        for val_path in tqdm.tqdm(val_paths):
            try:
                audio = Audio.from_file(val_path)
                result = model.transcribe(
                    audio.convert(16000, 2, 1).array_float.copy(),
                    initial_prompt=None,
                    condition_on_previous_text=False,
                )
                # result is: text, segment, lang
                # print(result["language"])
                whisper_outs.append(clean_text(result["text"]))
                whisper_results[val_path] = result
            except Exception as e:
                print(e)
                whisper_outs.append("")
                whisper_results[val_path] = {}
        # will take 1.5 mins
        default_wer = WordErrorRate(device="cuda")
        default_wer.update(whisper_outs, val_truth_text[: len(whisper_outs)])
        wer_value = default_wer.compute().item()
        print(validation_set, wer_value)
        log_dict[validation_set] = wer_value
        wandb.log(log_dict)
        with open(f"/home/tony/Data/Hoot/whipser/{validation_set}.json", "w") as fp:
            json.dump(whisper_results, fp)
        print("DONE!!!")