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

"""Create the SUC3-mini NER dataset and upload it to the HF Hub."""

import bz2
import io
import re

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 lxml import etree


def main() -> None:
    """Create the DaNE-mini NER dataset and uploads it to the HF Hub."""
    # Define download URLs
    url = "https://spraakbanken.gu.se/lb/resurser/meningsmangder/suc3.xml.bz2"

    # Download the data
    compressed_xml_data = requests.get(url).content

    # Decompress the data
    xml_data = bz2.decompress(compressed_xml_data)

    # Parse the XML data
    context = etree.iterparse(io.BytesIO(xml_data), events=("start", "end"))

    # Define a conversion dict to process the NER tags
    conversion_dict = dict(
        O="O",
        animal="MISC",
        event="MISC",
        inst="ORG",
        myth="MISC",
        other="MISC",
        person="PER",
        place="LOC",
        product="MISC",
        work="MISC",
    )

    # Initialise a `ner_tag` variable as the "no tag", `O`
    ner_tag = "O"

    # Iterate over the context
    records: list[dict[str, str | list[str]]] = list()
    tokens: list[str] = list()
    ner_tags: list[str] = list()
    for action, elt in context:
        # If the current element begins a name then set the `ner_tag` to the
        # corresponding `B` value
        if elt.tag == "name" and action == "start":
            ner_tag = f"B-{conversion_dict[elt.attrib['type']]}"

        # If the current element ends a name then reset the `ner_tag` to `O`
        elif elt.tag == "name" and action == "end":
            ner_tag = "O"

        # If the current element starts a word (i.e., `w`) then add the word to the
        # list of tokens, and add the current NER tag to the list of NER tags
        elif elt.tag == "w" and action == "start":
            if elt.text:
                tokens.append(elt.text)
                ner_tags.append(ner_tag)

        # If the current element ends a word then set the current NER tag to the `I`
        # version of the previous NER tag
        elif elt.tag == "w" and action == "end":
            if ner_tag.startswith("B-"):
                ner_tag = f"I-{ner_tag[2:]}"

        # If the current element ends a sentence then store all the data in the list of
        # records
        elif elt.tag == "sentence" and action == "end" and len(tokens) > 0:
            # Create document from the tokens
            doc = re.sub(" ([.,])", "\1", " ".join(tokens))

            # Double-check that there are the same number of tokens as NER tags
            assert len(tokens) == len(ner_tags)

            # Create the record and append it to the list of records
            record: dict[str, str | list[str]] = dict(
                doc=doc, tokens=tokens, ner_tags=ner_tags
            )
            records.append(record)

        # If the current element starts a sentence then reset the list of tokens and
        # NER tags, and set the current NER tag to `O`
        elif elt.tag == "sentence" and action == "start":
            tokens = list()
            ner_tags = list()
            ner_tag = "O"

    # Convert the records to a DataFrame
    df = pd.DataFrame.from_records(records)

    # 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),
        }
    )

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


if __name__ == "__main__":
    main()
