import os
import fire
import numpy as np
import pandas as pd
import shutil
import soundfile
import re

from slue_toolkit.prepare import data_utils
from slue_toolkit.prepare.create_dict import create_dict
from slue_toolkit.generic_utils import save_dct, write_to_file

splits = {"fine-tune", "dev", "test"}


def create_split(
    input_dir="/persist/data/voxpopuli", output_dir="/persist/data/slue-voxpopuli"
):
    os.makedirs(output_dir, exist_ok=True)
    for split in splits:
        print(f"processing {split}")
        split_dir = os.path.join(output_dir, split)
        os.makedirs(split_dir, exist_ok=True)
        metafile = os.path.join(input_dir, f"voxpopuli_asr_en_{split}_ner.tsv")
        shutil.copy(metafile, os.path.join(output_dir, f"{split}.tsv"))
        df = pd.read_csv(metafile, sep="\t")
        for id in df["id"].array:
            filename = os.path.join(
                input_dir, f"voxpopuli_asr_en/transcribed_data/en/{id[:4]}/{id}.ogg"
            )
            shutil.copy(filename, split_dir)


def create_manifest(
    data_dir="data/slue-voxpopuli",
    manifest_dir="manifest/slue-voxpopuli",
    is_blind=True,
):
    os.makedirs(manifest_dir, exist_ok=True)
    for split in splits:
        print(f"Processing slue-voxpopuli {split} set")
        if (split == "test") and is_blind:
            df = pd.read_csv(
                os.path.join(data_dir, f"slue-voxpopuli_{split}_blind.tsv"), sep="\t"
            )
        else:
            df = pd.read_csv(
                os.path.join(data_dir, f"slue-voxpopuli_{split}.tsv"), sep="\t"
            )

        with open(os.path.join(manifest_dir, f"{split}.tsv"), "w") as f:
            print(os.path.abspath(os.path.join(data_dir, split)), file=f)
            for uid in df["id"].array:
                frames = soundfile.info(
                    os.path.join(data_dir, split, f"{uid}.ogg")
                ).frames
                print(f"{uid}.ogg\t{frames}", file=f)

        for sub_dir_name in ["e2e_ner", "text_ner"]:
            os.makedirs(os.path.join(manifest_dir, sub_dir_name), exist_ok=True)
        if split != "test" or not is_blind:
            with open(os.path.join(manifest_dir, f"{split}.wrd"), "w") as f:
                for text in df["normalized_text"].array:
                    text = re.sub(r"[\.;?!]", "", text)
                    text = re.sub(r"\s+", " ", text)
                    print(text, file=f)

            with open(os.path.join(manifest_dir, f"{split}.ltr"), "w") as f:
                for text in df["normalized_text"].array:
                    text = re.sub(r"[\.;?!]", "", text)
                    text = re.sub(r"\s+", " ", text)
                    print(" ".join(text.replace(" ", "|")), file=f)

            # prepare NER files (for Fairseq and HuggingFace)
            for label_type in ["raw", "combined"]:
                wrd_fn = os.path.join(
                    manifest_dir, "e2e_ner", f"{split}.{label_type}.wrd"
                )
                ltr_fn = os.path.join(
                    manifest_dir, "e2e_ner", f"{split}.{label_type}.ltr"
                )
                tsv_fn = os.path.join(
                    manifest_dir, "text_ner", f"{split}.{label_type}.tsv"
                )
                with open(wrd_fn, "w") as f_wrd, open(ltr_fn, "w") as f_ltr, open(
                    tsv_fn, "w"
                ) as f_tsv:
                    for data_sample in df.iterrows():
                        entity_pair_str = data_utils.prep_text_ner_tsv(
                            data_sample[1].normalized_text,
                            data_sample[1].normalized_ner,
                            label_type,
                        )
                        print(entity_pair_str, file=f_tsv, end="")
                        wrd_str, ltr_str = data_utils.prep_e2e_ner_files(
                            entity_pair_str, label_type
                        )
                        print(wrd_str, file=f_wrd)
                        print(ltr_str, file=f_ltr)
        e2e_ner_tsv_fn = os.path.join(manifest_dir, f"e2e_ner/{split}.tsv")
        if not os.path.exists(e2e_ner_tsv_fn):
            os.symlink(f"{manifest_dir}/{split}.tsv", e2e_ner_tsv_fn)

    for label_type in ["raw", "combined"]:
        tag2id, id2tag, tag_lst_ordered = data_utils.prepare_tag_id_mapping(label_type)
        save_dct(
            os.path.join(manifest_dir, "text_ner", f"{label_type}_tag2id.json"), tag2id
        )
        save_dct(
            os.path.join(manifest_dir, "text_ner", f"{label_type}_id2tag.json"), id2tag
        )
        write_to_file(
            "\n".join(tag_lst_ordered),
            os.path.join(manifest_dir, "text_ner", f"{label_type}_tag_lst_ordered"),
        )
        for token_type in ["wrd", "ltr"]:
            create_dict(
                os.path.join(
                    manifest_dir, f"e2e_ner/fine-tune.{label_type}.{token_type}"
                ),
                os.path.join(
                    manifest_dir, f"e2e_ner/dict.{label_type}.{token_type}.txt"
                ),
            )


if __name__ == "__main__":
    fire.Fire()
