# /// 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 ScaLA datasets and upload them to the HF Hub."""

import random
import re
import warnings

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

from .constants import MAX_NUM_CHARS_IN_DOCUMENT, MIN_NUM_CHARS_IN_DOCUMENT  # noqa
from .load_ud_pos import (
    load_bgdt_pos,
    load_cadt_pos,
    load_csdt_pos,
    load_dadt_pos,
    load_dedt_pos,
    load_eldt_pos,
    load_endt_pos,
    load_esdt_pos,
    load_etdt_pos,
    load_fidt_pos,
    load_fodt_pos,
    load_frdt_pos,
    load_hrdt_pos,
    load_hudt_pos,
    load_isdt_pos,
    load_itdt_pos,
    load_ltdt_pos,
    load_lvdt_pos,
    load_nldt_pos,
    load_nodt_nb_pos,
    load_nodt_nn_pos,
    load_pldt_pos,
    load_ptdt_pos,
    load_rodt_pos,
    load_skdt_pos,
    load_sldt_pos,
    load_sqdt_pos,
    load_srdt_pos,
    load_svdt_pos,
    load_ukdt_pos,
)


def main() -> None:
    """Create the ScaLA datasets and upload them to the HF Hub."""
    # Set up the POS dataset loaders
    pos_datasets = {
        "da": load_dadt_pos,
        "nb": load_nodt_nb_pos,
        "nn": load_nodt_nn_pos,
        "sv": load_svdt_pos,
        "is": load_isdt_pos,
        "fo": load_fodt_pos,
        "de": load_dedt_pos,
        "nl": load_nldt_pos,
        "en": load_endt_pos,
        "fr": load_frdt_pos,
        "it": load_itdt_pos,
        "es": load_esdt_pos,
        "fi": load_fidt_pos,
        "pt": load_ptdt_pos,
        "lv": load_lvdt_pos,
        "et": load_etdt_pos,
        "pl": load_pldt_pos,
        "lt": load_ltdt_pos,
        "cs": load_csdt_pos,
        "sk": load_skdt_pos,
        "uk": load_ukdt_pos,
        "el": load_eldt_pos,
        "bg": load_bgdt_pos,
        "sr": load_srdt_pos,
        "sl": load_sldt_pos,
        "hr": load_hrdt_pos,
        "hu": load_hudt_pos,
        "ro": load_rodt_pos,
        "ca": load_cadt_pos,
        "sq": load_sqdt_pos,
    }

    # Set up the progress bar and iterate over the languages
    with tqdm(pos_datasets.items(), desc="Creating ScaLA datasets") as pbar:
        for lang, fn in pbar:
            # Update the progress bar description
            pbar.set_description(f"Creating ScaLA datasets - {lang}")

            # Load the POS dataset
            pos_dataset = fn()

            # Merge the DDT POS dataframes to a single dataframe, with columns `ids`,
            # `tokens`, `doc` and `pos_tags`
            df = pd.concat(pos_dataset.values(), ignore_index=True)
            assert isinstance(df, pd.DataFrame)

            # Drop the duplicates
            df = df.drop_duplicates(subset="doc").reset_index(drop=True)
            assert isinstance(df, pd.DataFrame)

            # Remove samples with five or fewer tokens
            df = df[df.tokens.map(lambda lst: len(lst) > 5)]
            assert isinstance(df, pd.DataFrame)

            # Remove samples with five or fewer distinct POS tags
            df = df[df.pos_tags.map(lambda lst: len(set(lst)) > 5)]
            assert isinstance(df, pd.DataFrame)

            # Remove samples with an odd number of quotes
            df = df[df.doc.map(lambda doc: doc.count('"') % 2 == 0)]
            assert isinstance(df, pd.DataFrame)

            # Remove samples which starts with punctuation
            df = df[df.pos_tags.map(lambda lst: lst[0] not in ["PUNCT", "SYM"])]
            assert isinstance(df, pd.DataFrame)

            # Remove samples containing more than one '=' character, as this is used to
            # indicate a tag
            df = df[df.doc.map(lambda doc: doc.count("=") <= 1)]
            assert isinstance(df, pd.DataFrame)

            # Remove samples containing 'SLUTORD', as this is used to indicate a tag
            df = df[~df.doc.str.contains("SLUTORD")]
            assert isinstance(df, pd.DataFrame)

            # Create a training, validation, test and small training set. Note that we
            # will corrupt the data, so this is only half the size of the final
            # datasets. In the case where the dataframe does not contain enough samples
            # for all the splits, we keep halving the test size until we have enough
            # samples.
            test_size = 1024
            while test_size >= 128:
                try:
                    train_df, val_df, test_df, full_train_df = make_splits(
                        df=df, train_size=512, val_size=128, test_size=test_size
                    )
                    break
                except ValueError:
                    test_size //= 2
            else:
                train_size = 64
                val_size = 32
                test_size = len(df) - train_size - val_size
                assert test_size > 0, "Not enough samples to create the splits."
                train_df, val_df, test_df, full_train_df = make_splits(
                    df=df, train_size=train_size, val_size=val_size, test_size=test_size
                )

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

            # Only work with samples where the document is not very large or small We do
            # it after we have made the splits to ensure that the dataset is minimally
            # affected.
            new_train_df = train_df.copy()
            new_train_df["text_len"] = new_train_df.doc.str.len()
            new_train_df = new_train_df.query(
                "text_len >= @MIN_NUM_CHARS_IN_DOCUMENT"
            ).query("text_len <= @MAX_NUM_CHARS_IN_DOCUMENT")
            new_val_df = val_df.copy()
            new_val_df["text_len"] = new_val_df.doc.str.len()
            new_val_df = new_val_df.query(
                "text_len >= @MIN_NUM_CHARS_IN_DOCUMENT"
            ).query("text_len <= @MAX_NUM_CHARS_IN_DOCUMENT")
            new_test_df = test_df.copy()
            new_test_df["text_len"] = new_test_df.doc.str.len()
            new_test_df = new_test_df.query(
                "text_len >= @MIN_NUM_CHARS_IN_DOCUMENT"
            ).query("text_len <= @MAX_NUM_CHARS_IN_DOCUMENT")
            new_full_train_df = full_train_df.copy()
            new_full_train_df["text_len"] = new_full_train_df.doc.str.len()
            new_full_train_df = new_full_train_df.query(
                "text_len >= @MIN_NUM_CHARS_IN_DOCUMENT"
            ).query("text_len <= @MAX_NUM_CHARS_IN_DOCUMENT")

            # Add the corrupted data and turn the dataframes into Hugging Face Dataset
            # objects
            train = prepare_df(new_train_df, split="train")
            val = prepare_df(new_val_df, split="val")
            test = prepare_df(new_test_df, split="test")
            full_train = prepare_df(new_full_train_df, split="train")

            # Collect datasets in a dataset dictionary
            dataset = DatasetDict(
                {"train": train, "val": val, "test": test, "full_train": full_train}
            )

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


def join_tokens(tokens: list[str]) -> str:
    """Joins a list of tokens into a string.

    Args:
        tokens:
            The list of tokens to join.

    Returns:
        The joined string.
    """
    # Form document
    doc = " ".join(tokens)

    # Remove whitespace around punctuation
    doc = (
        doc.replace(" .", ".")
        .replace(" ,", ",")
        .replace(" ;", ";")
        .replace(" :", ":")
        .replace("( ", "(")
        .replace(" )", ")")
        .replace("[ ", "[")
        .replace(" ]", "]")
        .replace("{ ", "{")
        .replace(" }", "}")
        .replace(" ?", "?")
        .replace(" !", "!")
    )

    # Remove whitespace around quotes
    if doc.count('"') % 2 == 0:
        doc = re.sub('" ([^"]*) "', '"\\1"', doc)

    # Return the document
    return doc


def delete(tokens: list[str], pos_tags: list[str]) -> str | None:
    """Delete a random token from a list of tokens.

    The POS tags are used to prevent deletion of a token which does not make the
    resulting sentence grammatically incorrect, such as removing an adjective or an
    adverb.

    Args:
        tokens:
            The list of tokens to delete from.
        pos_tags:
            The list of POS tags for the tokens.

    Returns:
        The deleted token, or None if no token could be deleted.
    """
    # Copy the token list
    new_tokens = tokens.copy()

    # Get candidate indices to remove. We do not remove adjectives, adverbs,
    # punctuation, determiners or numbers, as the resulting sentence will probably
    # still be grammatically correct. Further, we do not remove nouns or proper nouns
    # if they have another noun or proper noun as neighbour, as that usually does not
    # make the sentence incorrect either.
    indices = [
        idx
        for idx, pos_tag in enumerate(pos_tags)
        if pos_tag not in ["ADJ", "ADV", "PUNCT", "SYM", "DET", "NUM", "PART"]
        and (
            pos_tag not in ["NOUN", "PROPN"]
            or (
                (idx == 0 or pos_tags[idx - 1] not in ["NOUN", "PROPN"])
                and (
                    idx == len(new_tokens) - 1
                    or pos_tags[idx + 1] not in ["NOUN", "PROPN"]
                )
            )
        )
    ]

    # If there are no candidates then return None
    if len(indices) == 0:
        return None

    # Get the random index
    rnd_idx = random.choice(indices)

    # Delete the token at the index
    new_tokens.pop(rnd_idx)

    # Join up the new tokens and return the string
    return join_tokens(new_tokens)


def flip_neighbours(tokens: list[str], pos_tags: list[str]) -> str | None:
    """Flip a pair of neighbouring tokens.

    The POS tags are used to prevent flipping of tokens which does not make the
    resulting sentence grammatically incorrect, such as flipping two adjectives.

    Args:
        tokens:
            The list of tokens to flip.
        pos_tags:
            The list of POS tags for the tokens.

    Returns:
        The flipped string, or None if no flip was possible.
    """
    # Copy the token list
    new_tokens = tokens.copy()

    # Collect all indices that are proper words, and which has a neighbour which is
    # also a proper word as well as having a different POS tag
    indices = [
        idx for idx, pos_tag in enumerate(pos_tags) if pos_tag not in ["PUNCT", "SYM"]
    ]
    indices = [
        idx
        for idx in indices
        if (idx + 1 in indices and pos_tags[idx] != pos_tags[idx + 1])
        or (idx - 1 in indices and pos_tags[idx] != pos_tags[idx - 1])
    ]

    # If there are fewer than two relevant tokens then return None
    if len(indices) < 2:
        return None

    # Get the first random index
    rnd_fst_idx = random.choice(indices)

    # Get the second (neighbouring) index
    if rnd_fst_idx == 0:
        rnd_snd_idx = rnd_fst_idx + 1
    elif rnd_fst_idx == len(tokens) - 1:
        rnd_snd_idx = rnd_fst_idx - 1
    elif (
        pos_tags[rnd_fst_idx + 1] in ["PUNCT", "SYM"]
        or pos_tags[rnd_fst_idx] == pos_tags[rnd_fst_idx + 1]
        or {pos_tags[rnd_fst_idx], pos_tags[rnd_fst_idx + 1]} == {"PRON", "AUX"}
    ):
        rnd_snd_idx = rnd_fst_idx - 1
    elif (
        pos_tags[rnd_fst_idx - 1] in ["PUNCT", "SYM"]
        or pos_tags[rnd_fst_idx] == pos_tags[rnd_fst_idx - 1]
        or {pos_tags[rnd_fst_idx], pos_tags[rnd_fst_idx + 1]} == {"PRON", "AUX"}
    ):
        rnd_snd_idx = rnd_fst_idx + 1
    elif random.random() > 0.5:
        rnd_snd_idx = rnd_fst_idx - 1
    else:
        rnd_snd_idx = rnd_fst_idx + 1

    # Flip the two indices
    new_tokens[rnd_fst_idx] = tokens[rnd_snd_idx]
    new_tokens[rnd_snd_idx] = tokens[rnd_fst_idx]

    # If we flipped the first character, then ensure that the new first character is
    # title-cased and the second character is of lower case. We only do this if they
    # are not upper cased, however.
    if rnd_fst_idx == 0 or rnd_snd_idx == 0:
        if new_tokens[0] != new_tokens[0].upper():
            new_tokens[0] = new_tokens[0].title()
        if new_tokens[1] != new_tokens[1].upper():
            new_tokens[1] = new_tokens[1].lower()

    # Join up the new tokens and return the string
    return join_tokens(new_tokens)


def corrupt(
    tokens: list[str], pos_tags: list[str], num_corruptions: int = 1
) -> list[tuple[str, str]]:
    """Corrupt a list of tokens.

    This randomly either flips two neighbouring tokens or deletes a random token.

    Args:
        tokens:
            The list of tokens to corrupt.
        pos_tags:
            The list of POS tags for the tokens.
        num_corruptions:
            The number of corruptions to perform. Defaults to 1.

    Returns:
        The list of (corrupted_string, corruption_type)
    """
    # Define the list of corruptions
    corruptions: list[tuple[str, str]] = list()

    # Continue until we have achieved the desired number of corruptions
    while len(corruptions) < num_corruptions:
        # Choose which corruption to perform, at random
        corruption_fn = random.choice([flip_neighbours, delete])

        # Corrupt the tokens
        corruption = corruption_fn(tokens, pos_tags)

        # If the corruption succeeded, and that we haven't already performed the same
        # corruption, then add the corruption to the list of corruptions
        if corruption not in corruptions and corruption is not None:
            corruptions.append((corruption, corruption_fn.__name__))

    # Return the list of corruptions
    return corruptions


def prepare_df(df: pd.DataFrame, split: str) -> Dataset:
    """Prepare a dataframe by adding an equal number of corruptions to it.

    Args:
        df:
            The dataframe to prepare.
        split:
            The split to prepare the dataframe for.

    Returns:
        The prepared dataset.
    """
    # Reset the index of the dataframe
    df.reset_index(drop=True, inplace=True)

    # Get the corrupted strings
    corrupted_list = [
        corrupt(tokens=tokens, pos_tags=pos_tags, num_corruptions=1)[0]
        for tokens, pos_tags in zip(df.tokens, df.pos_tags)
    ]

    # Add the corrupted strings to the dataframe
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", category=SettingWithCopyWarning)
        df["corrupted"] = [tup[0] for tup in corrupted_list]
        df["corruption_type"] = [tup[1] for tup in corrupted_list]

    # Restructure the dataframe to have columns `text`, `corruption_type` and `label`,
    # with one sample per row
    df = pd.concat(
        [
            pd.DataFrame(
                dict(
                    text=df.tokens.map(join_tokens).tolist(),
                    corruption_type=[None for _ in range(len(df))],
                    label=["correct" for _ in range(len(df))],
                )
            ),
            pd.DataFrame(
                dict(
                    text=df.corrupted.explode().tolist(),
                    corruption_type=df.corruption_type.explode().tolist(),
                    label=["incorrect" for _ in range(len(df))],
                )
            ),
        ]
    )

    # Shuffle the dataframe
    df = df.sample(frac=1.0, random_state=4242).reset_index(drop=True)

    # Convert the dataframe to a Hugging Face Dataset and return it
    return Dataset.from_pandas(df, split=split)  # type: ignore[bad-argument-type]


def make_splits(
    df: pd.DataFrame, train_size: int, val_size: int, test_size: int
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Make the splits.

    Args:
        df:
            The dataframe to make the splits from.
        train_size:
            The size of the training split.
        val_size:
            The size of the validation split.
        test_size:
            The size of the test split.

    Returns:
        The train, validation, test, and full training splits.
    """
    val_df = df.sample(n=val_size, random_state=4242)
    df_filtered = df[~df.index.isin(val_df.index)]
    test_df = df_filtered.sample(n=test_size, 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=train_size, random_state=4242)
    return train_df, val_df, test_df, full_train_df


if __name__ == "__main__":
    main()
