# /// script
# requires-python = ">=3.10,<4.0"
# dependencies = [
#     "datasets==3.5.0",
#     "huggingface-hub==0.24.0",
#     "pandas==2.2.0",
#     "requests==2.32.3",
#     "tqdm==4.67.1",
# ]
# ///

"""Create the NorNE-mini NER datasets and upload them to the HF Hub."""

import re
import typing as t
from collections import defaultdict

import pandas as pd
import requests
from datasets import Split
from datasets.arrow_dataset import Dataset
from datasets.dataset_dict import DatasetDict
from huggingface_hub.hf_api import HfApi
from tqdm.auto import tqdm


def main() -> None:
    """Create the NorNE-mini NER datasets and uploads them to the HF Hub."""
    # Define the base download URLs
    base_urls = dict(
        nb="https://github.com/UniversalDependencies/UD_Norwegian-Bokmaal/raw/master/"
        "no_bokmaal-ud-{}.conllu",
        nn="https://github.com/UniversalDependencies/UD_Norwegian-Nynorsk/raw/master/"
        "no_nynorsk-ud-{}.conllu",
    )

    # Define conversion dict to be used with the NER tags
    ner_conversion_dict = {
        "O": "O",
        "B-LOC": "B-LOC",
        "I-LOC": "I-LOC",
        "B-PER": "B-PER",
        "I-PER": "I-PER",
        "B-ORG": "B-ORG",
        "I-ORG": "I-ORG",
        "B-MISC": "B-MISC",
        "I-MISC": "I-MISC",
        "B-GPE_LOC": "B-LOC",
        "I-GPE_LOC": "I-LOC",
        "B-GPE_ORG": "B-ORG",
        "I-GPE_ORG": "I-ORG",
        "B-PROD": "B-MISC",
        "I-PROD": "I-MISC",
        "B-DRV": "B-MISC",
        "I-DRV": "I-MISC",
        "B-EVT": "B-MISC",
        "I-EVT": "I-MISC",
    }

    # Set up the progress bar and iterate over the languages
    with tqdm(["nb", "nn"], desc="Creating NorNE-mini datasets") as pbar:
        for lang in pbar:
            # Define download URLs
            train_url = base_urls[lang].format("train")
            val_url = base_urls[lang].format("dev")
            test_url = base_urls[lang].format("test")

            # Download the data
            data = dict(
                train=requests.get(train_url).text.split("\n"),
                val=requests.get(val_url).text.split("\n"),
                test=requests.get(test_url).text.split("\n"),
            )

            # Update the progress bar description
            pbar.set_description(f"Creating NorNE datasets - {lang}")

            # Iterate over the data splits
            dfs = dict()
            for split, lines in data.items():
                # Initialise the records, data dictionary and document
                records = list()
                data_dict: t.Dict[str, t.List[t.Union[int, str]]] = defaultdict(list)
                doc = ""

                # Iterate over the data for the given split
                for line in lines:
                    # If we are at the first line of an entry then extract the document
                    if line.startswith("# text = "):
                        doc = re.sub("# text = ", "", line)

                    # Otherwise, if the line is a comment then ignore it
                    elif line.startswith("#"):
                        continue

                    # Otherwise, if we have reached the end of an entry then store it
                    # to the list of records and reset the data dictionary and document
                    elif line == "":
                        merged_data_dict: t.Dict[
                            str, t.Union[str, t.List[t.Union[int, str]]]
                        ]
                        if len(data_dict["tokens"]) > 0:
                            merged_data_dict = {**data_dict, "doc": doc}
                            records.append(merged_data_dict)
                        data_dict = defaultdict(list)
                        doc = ""

                    # Otherwise we are in the middle of an entry which is not a
                    # comment, so we extract the data from the line and store it in the
                    # data dictionary
                    else:
                        data_tup = line.split("\t")
                        ner_tag = data_tup[9].replace("name=", "").split("|")[-1]
                        data_dict["ids"].append(data_tup[0])
                        data_dict["tokens"].append(data_tup[1])
                        data_dict["ner_tags"].append(ner_conversion_dict[ner_tag])

                # Convert the records to a dataframe and store it
                dfs[split] = pd.DataFrame.from_records(records)

            # Merge the dataframes
            df = pd.concat(dfs.values()).reset_index(drop=True)

            # Remove duplicates
            df = df.drop_duplicates().reset_index(drop=True)

            # Create new splits
            val_df = df.sample(n=256, random_state=4242)
            df_filtered = df[~df.index.isin(val_df.index)]
            test_df = df_filtered.sample(n=2048, random_state=4242)
            full_train_df = df_filtered[~df_filtered.index.isin(test_df.index)]
            assert isinstance(full_train_df, pd.DataFrame)
            train_df = full_train_df.sample(n=1024, random_state=4242)

            assert isinstance(train_df, pd.DataFrame)
            assert isinstance(val_df, pd.DataFrame)
            assert isinstance(test_df, pd.DataFrame)

            # Collect datasets in a dataset dictionary
            dataset = DatasetDict(
                {
                    "train": Dataset.from_pandas(train_df, split=Split.TRAIN),
                    "val": Dataset.from_pandas(val_df, split=Split.VALIDATION),
                    "test": Dataset.from_pandas(test_df, split=Split.TEST),
                    "full_train": Dataset.from_pandas(
                        full_train_df,
                        split="full_train",  # type: ignore[bad-argument-type]
                    ),
                }
            )

            # Push the dataset to the Hugging Face Hub
            dataset_id = f"EuroEval/norne-{lang}-mini"
            HfApi().delete_repo(dataset_id, repo_type="dataset", missing_ok=True)
            dataset.push_to_hub(dataset_id, private=True)


if __name__ == "__main__":
    main()
