"""
# original en version (v1)
CUDA_VISIBLE_DEVICES=0 python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/Hoot/default_model.pt \
    --tokenizer_filepath=/home/tony/Data/Hoot/tokenizer_v1024_default.model \
    --wandb_run_name=hoot_v1_en

#  hoot filtered lyrics (v2)
python hoot_evaluation.py \
    --use_prior=True \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2023-11-22_07-04-38/50k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/full/tokenizer_spe_bpe_v20480/tokenizer.model \
    --wandb_run_name=hoot_multi_20k_vocab_hooted_fine_tune_50k   

# v3 tries t4
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2023-12-30_15-08-02/190k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/full/tokenizer_spe_bpe_v20480/tokenizer.model \
    --wandb_run_name=hoot_multi_v3_t4_190k 

# v3 tries t5 -- different tokenzier
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2024-01-02_21-53-24/180k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/v3/tokenizer_spe_bpe_v10240/tokenizer.model \
    --wandb_run_name=hoot_multi_v3_t5_180k 
    
# v3 tries t7
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2024-01-08_00-43-32/175k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/full/tokenizer_spe_bpe_v20480/tokenizer.model \
    --wandb_run_name=hoot_multi_v3_t7_175k_ckpt.pt  

# v3 tries t8 -- shorter training, faster lr
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2024-01-10_05-56-59/70k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/full/tokenizer_spe_bpe_v20480/tokenizer.model \
    --wandb_run_name=hoot_multi_v3_t8_70k  

# v3 tries t9 
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2024-01-12_22-51-29/170k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/full/tokenizer_spe_bpe_v20480/tokenizer.model \
    --wandb_run_name=hoot_multi_v3_t9_170k  

# v3 tries t10
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2024-01-15_18-58-13/60k_ckpt.pt \
    --tokenizer_filepath=/home/tony/Work/tony/hoot/tokenizers/full/tokenizer_spe_bpe_v20480/tokenizer.model \
    --wandb_run_name=hoot_multi_v3_t10_60k

# v3 validation
python hoot_evaluation.py \
    --checkpoint_filepath=/app/suno/models/hoot_v3.pt \
    --tokenizer_filepath=/app/suno/models/hoot_v3_tokenizer.model \
    --wandb_run_name=hoot_v3

# v4 validation
python hoot_evaluation.py \
    --checkpoint_filepath=/home/tony/Data/checkpoints/hoot/2025-01-15_03-43-37/25k_ckpt.pt \
    --tokenizer_filepath=/app/suno/models/hoot_v3_tokenizer.model \
    --wandb_run_name=hoot_v4_25k
"""

import os
import json
import argparse
from suno_utils.tasks.hoot import load_model_list, encode_filepaths
import wandb
import tqdm
from torcheval.metrics import WordErrorRate
import regex
import copy
import numpy as np
import time
import datetime
from suno_utils.utils.lyrics import remove_speakers
from suno_utils.utils.text import normalize_whitespace


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


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--checkpoint_filepath", type=str)
    parser.add_argument("--tokenizer_filepath", type=str)
    parser.add_argument("--wandb_run_name", type=str)
    parser.add_argument("--use_prior", type=bool, default=False)
    args = parser.parse_args()
    return args


LOG_WANDB = True
if __name__ == "__main__":
    avialbe_device = f"cuda:{os.environ['CUDA_VISIBLE_DEVICES']}"
    print(avialbe_device)
    input_args = parse_args()
    model_list = load_model_list(
        checkpoint_filepath=input_args.checkpoint_filepath,
        tokenizer_filepath=input_args.tokenizer_filepath,
    )
    model = model_list[0]["model"]
    tokenizer = model_list[0]["tokenizer"]
    if LOG_WANDB:
        wandb.init(
            project="hoot_evals",
            name=f"{input_args.wandb_run_name}{'_prior' if input_args.use_prior else ''}",
        )
    validation_datasets = {
        "v4_en": "/home/tony/Data/Hoot/v4_en_validation_set.json",
        # "multi": "/home/tony/Data/Hoot/all_test_manifest.json",
        # "multi_balanced": "/home/tony/Data/Hoot/multi_balanced_test_manifest.json",
        "v4_multi": "/home/tony/Data/Hoot/v4_multi_validation_set.json",
    }
    for validation_set, validation_path in validation_datasets.items():
        log_dict = {}
        val_paths = []
        val_truth_text = []
        default_hoot_outs = []
        with open(validation_path, "r") as fp:
            for l in tqdm.tqdm(fp):
                val_meta = json.loads(l)
                val_paths.append(val_meta["audio_filepath"])
                val_truth_text.append(val_meta["text"])
                # if input_args.use_prior:
                #     out = encode_filepaths([val_meta["audio_filepath"]], return_logits=True)[0]
                #     pred_text = tokenizer.decode_logits(out, prior_text=val_meta["text"].lower())
                #     default_hoot_outs.append(clean_text(pred_text))
        # will take 1.5 mins
        if not input_args.use_prior:
            print(
                f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, start encoding the validation set"
            )
            default_hoot_outs = []
            chunk_size = 100
            for start_idx in tqdm.tqdm(range(0, len(val_paths), chunk_size)):
                val_path_chunk = val_paths[start_idx : start_idx + chunk_size]
                default_hoot_outs.extend(
                    encode_filepaths(val_path_chunk, batch_size=128)
                )
            print(
                f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, finished encoding the validation set"
            )
        else:
            print(
                f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, start encoding the validation set"
            )
            default_hoot_outs = []
            chunk_size = 100
            for i in tqdm.tqdm(range(0, len(val_paths), chunk_size)):
                val_path_chunk = val_paths[i : i + chunk_size]
                val_truth_text_chunk = [
                    x.lower() for x in val_truth_text[i : i + chunk_size]
                ]
                default_hoot_outs.extend(
                    encode_filepaths(
                        val_path_chunk, batch_size=128, prior_texts=val_truth_text_chunk
                    )
                )
            print(
                f"{datetime.datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}, finished encoding the validation set"
            )

        default_wer = WordErrorRate(device="cuda")
        default_wer.update(default_hoot_outs, val_truth_text)
        wer_value = default_wer.compute().item()
        print(validation_set, wer_value)
        log_dict[validation_set] = wer_value
        if LOG_WANDB:
            wandb.log(log_dict)
    print("DONE!!!")
