# /// script
# requires-python = ">=3.10,<4.0"
# dependencies = [
#     "pandas==2.2.0",
#     "requests==2.32.3",
# ]
# ///

"""Load the part-of-speech part of a Universal Dependencies treebank."""

import collections.abc as c
import logging
import re
import typing as t
from collections import defaultdict
from functools import partial
from urllib.parse import urlparse

import pandas as pd
import requests

logging.basicConfig(format="%(asctime)s ⋅ %(message)s", level=logging.INFO)
logger = logging.getLogger("load_ud_pos")


def load_dadt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Danish Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Danish-DDT/raw/master/"
        "da_ddt-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_ptdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Portuguese Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Portuguese-Bosque/raw/master/"
        "pt_bosque-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(
        train_url=train_url,
        val_url=val_url,
        test_url=test_url,
        filter_source="CETEMPúblico",
    )


def load_fodt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Faroese Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Faroese-FarPaHC/raw/master/"
        "fo_farpahc-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_isdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Icelandic Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Icelandic-Modern/raw/master/"
        "is_modern-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_nodt_nb_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Norwegian Bokmål Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Norwegian-Bokmaal/raw/master/"
        "no_bokmaal-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_nodt_nn_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Norwegian Nynorsk Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Norwegian-Nynorsk/raw/master/"
        "no_nynorsk-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_svdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Swedish Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://github.com/UniversalDependencies/UD_Swedish-Talbanken/raw/master/"
        "sv_talbanken-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    # Define document processing function
    def process_document(doc: str) -> str:
        doc = (
            doc.replace(" s k", " s.k.")
            .replace("S k", "S.k.")
            .replace(" bl a", " bl.a.")
            .replace("Bl a", "Bl.a.")
            .replace(" t o m", " t.o.m.")
            .replace("T o m", "T.o.m.")
            .replace(" fr o m", " fr.o.m.")
            .replace("Fr o m", "Fr.o.m.")
            .replace(" o s v", " o.s.v.")
            .replace("O s v", "O.s.v.")
            .replace(" d v s", " d.v.s.")
            .replace("D v s", "D.v.s.")
            .replace(" m fl", " m.fl.")
            .replace("M fl", "M.fl.")
            .replace(" t ex", " t.ex.")
            .replace("T ex", "T.ex.")
            .replace(" f n", " f.n.")
            .replace("F n", "F.n.")
        )
        return doc

    return load_ud_pos(
        train_url=train_url,
        val_url=val_url,
        test_url=test_url,
        doc_process_fn=process_document,
    )


def load_dedt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the German Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_German-GSD/master/"
        "de_gsd-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_nldt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Dutch Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Dutch-Alpino/"
        "master/nl_alpino-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_endt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the English Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_English-GUM/"
        "master/en_gum-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_frdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the French Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_French-GSD/"
        "master/fr_gsd-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_itdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Italian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Italian-ISDT/"
        "master/it_isdt-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_esdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Spanish Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Spanish-AnCora"
        "/refs/heads/master/es_ancora-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_fidt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Finnish Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Finnish-TDT/"
        "master/fi_tdt-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_lvdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Latvian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Latvian-LVTB/"
        "master/lv_lvtb-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_etdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Estonian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Estonian-EDT/"
        "refs/heads/master/et_edt-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_pldt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Polish Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Polish-PDB/"
        "refs/heads/master/pl_pdb-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_ltdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Lithuanian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Lithuanian-ALKSNIS/"
        "master/lt_alksnis-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_csdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Czech Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Czech-CAC/refs/"
        "heads/master/cs_cac-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_skdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Slovak Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Slovak-SNK/refs"
        "/heads/master/sk_snk-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_ukdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Ukrainian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/"
        "UD_Ukrainian-ParlaMint/refs/heads/master/uk_parlamint-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_eldt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Greek Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Greek-GUD/refs"
        "/heads/master/el_gud-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_bgdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Bulgarian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Bulgarian-BTB/refs"
        "/heads/master/bg_btb-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_srdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Serbian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Serbian-SET/refs"
        "/heads/master/sr_set-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_sldt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Slovene Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Slovene-SSJ/refs"
        "/heads/master/sl_ssj-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_hrdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Croatian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Croatian-SET/refs"
        "/heads/master/hr_set-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_hudt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Hungarian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Hungarian-Szeged/refs"
        "/heads/master/hu_szeged-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_rodt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Romanian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Romanian-RRT/refs"
        "/heads/master/ro_rrt-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_cadt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Catalan Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Catalan-AnCora/refs"
        "/heads/master/ca_ancora-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def load_sqdt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Albanian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Albanian-STAF/refs"
        "/heads/master/sq_staf-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    df_dict = load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)

    # Load additional TSA test data and concatenate with the existing test split
    tsa_test_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Albanian-TSA/refs"
        "/heads/master/sq_tsa-ud-test.conllu"
    )
    tsa_test_lines = _load_file_or_url(url_or_path=tsa_test_url)
    tsa_test_df = _load_split(lines=tsa_test_lines)

    df_dict["test"] = pd.concat([df_dict["test"], tsa_test_df], ignore_index=True)

    return df_dict


def load_bedt_pos() -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of the Belarusian Dependency Treebank.

    Returns:
        The dataframes, stored in the keys `train`, `val`, and `test`.
    """
    # Define download URLs
    base_url = (
        "https://raw.githubusercontent.com/UniversalDependencies/UD_Belarusian-HSE/refs"
        "/heads/master/be_hse-ud-{}.conllu"
    )
    train_url = base_url.format("train")
    val_url = base_url.format("dev")
    test_url = base_url.format("test")

    return load_ud_pos(train_url=train_url, val_url=val_url, test_url=test_url)


def _load_file_or_url(url_or_path: str) -> list[str]:
    """Load a file from a URL or local path.

    Args:
        url_or_path: The URL or local path to load.

    Returns:
        The list of strings, one per line.
    """
    parsed = urlparse(url_or_path)
    if parsed.scheme.lower() in ("http", "https"):
        return requests.get(url_or_path).text.split("\n")
    else:
        with open(url_or_path, "r") as f:
            logger.warning(f"Loading data from local file: {url_or_path}")
            return [line.strip() for line in f.readlines()]


def _filter_token_range(data_dict: dict[str, list]) -> dict[str, list]:
    """Filter out tokens that belong to ranges in UD source files.

    Tokens that span more than one position are not supported by create_scala's
    prepare_df logic.

    Example files:

    - tests/test_scripts/test_create_scala/test_data/de_gsd-ud-train.conllu.adp_det
    - tests/test_scripts/test_create_scala/test_data/pl_pdb-ud-train.conllu.aux_clitic_*

    Args:
        data_dict: The input data dictionary.

    Returns:
        The filtered data dictionary. Its format is identical to the input.
    """
    output: t.DefaultDict[str, list] = defaultdict(list)

    range_start: int = 0
    range_end: int = 0
    # Track the index where we added the multi-word token so we can update its POS tag
    mwt_index: int | None = None

    # Pattern used to detect merged IDs
    merged_ids_pattern = re.compile(r"(\d+)-(\d+)", re.I)

    for i in range(len(data_dict["ids"])):
        match = merged_ids_pattern.match(data_dict["ids"][i])
        if match is not None:
            output["ids"].append(data_dict["ids"][i])
            output["tokens"].append(data_dict["tokens"][i])
            output["pos_tags"].append(data_dict["pos_tags"][i])
            range_start = int(match.group(1))
            range_end = int(match.group(2))
            # Remember where we added this multi-word token
            mwt_index = len(output["pos_tags"]) - 1
        else:
            token_id = int(data_dict["ids"][i].split(".")[0])
            if token_id >= range_start and token_id <= range_end:
                # If this is the first token in the range and we have a multi-word token
                # to update, copy the POS tag from this token to the multi-word token
                if mwt_index is not None and token_id == range_start:
                    output["pos_tags"][mwt_index] = data_dict["pos_tags"][i]
                    mwt_index = None  # Reset so we only update once per range
                # Skip token if in range
                continue
            else:
                output["ids"].append(data_dict["ids"][i])
                output["tokens"].append(data_dict["tokens"][i])
                output["pos_tags"].append(data_dict["pos_tags"][i])
                # Reset mwt_index when we move past the range
                mwt_index = None

    return output


def _load_split(
    *,
    lines: list[str],
    filter_source: str | None = None,
    doc_process_fn: c.Callable[[str], str] = lambda x: x,
) -> pd.DataFrame:
    """Load single split of the POS part of a Universal Dependencies treebank.

    Args:
        lines: The lines of the file to process.
        filter_source:
            If not `None`, only include entries with this source in the dataset.
        doc_process_fn:
            A function to apply to each document before parsing it.

    Returns:
        The dataframe for the given split.
    """
    # Initialise the records, data dictionary and document
    records: list[dict] = []
    data_dict: dict[str, list[list[int | str] | str]] = defaultdict(list)

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

            # Process the document if needed
            doc = doc_process_fn(doc)

        elif line.startswith("# source = "):
            source = line.removeprefix("# source = ").strip()

        # 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 == "" or line_idx == len(lines) - 1:
            if len(data_dict["tokens"]) > 0:
                if filter_source is None or filter_source in source:
                    merged_data_dict: dict[str, str | list[int | str]]
                    merged_data_dict = {**_filter_token_range(data_dict), "doc": doc}
                    records.append(merged_data_dict)
            data_dict = defaultdict(list)
            doc = ""
            source = ""

        # 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")
            data_dict["ids"].append(data_tup[0])
            data_dict["tokens"].append(data_tup[1])
            data_dict["pos_tags"].append(data_tup[3])

    # Convert the records to a dataframe
    return pd.DataFrame.from_records(records)


def load_ud_pos(
    train_url: str,
    val_url: str,
    test_url: str,
    doc_process_fn: c.Callable[[str], str] = lambda x: x,
    filter_source: str | None = None,
) -> dict[str, pd.DataFrame]:
    """Load the part-of-speech part of a Universal Dependencies treebank.

    Args:
        train_url:
            The URL of the training data.
        val_url:
            The URL of the validation data.
        test_url:
            The URL of the test data.
        doc_process_fn:
            A function to apply to each document before parsing it.
        filter_source:
            If not `None`, only include entries with this source in the dataset.

    Returns:
        The dataframes, stored in the keys `train`, `val` and `test`.
    """
    if filter_source is not None:
        logger.warning(
            f"Warning: Filtering dataset to include only entries with {filter_source=}"
        )

    # Load the lines of the splits
    train_lines = _load_file_or_url(url_or_path=train_url)
    val_lines = _load_file_or_url(url_or_path=val_url)
    test_lines = _load_file_or_url(url_or_path=test_url)

    # Load the splits
    split_loader = partial(
        _load_split, filter_source=filter_source, doc_process_fn=doc_process_fn
    )
    dfs = dict(
        train=split_loader(lines=train_lines),
        val=split_loader(lines=val_lines),
        test=split_loader(lines=test_lines),
    )

    # Return the dictionary of dataframes
    return dfs
