import pyrootutils

root = pyrootutils.setup_root(
    search_from=__file__,
    indicator=[".project-root"],
    pythonpath=True,
    dotenv=True,
)

import argparse
import logging
import os
from typing import List, Optional, TypeVar

from pie_datasets import load_dataset
from pie_datasets.builders.brat import BratDocument, BratDocumentWithMergedSpans
from pytorch_ie.documents import (
    TextBasedDocument,
    TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions,
    TextDocumentWithLabeledPartitions,
    TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions,
)

from src.document.processing import replace_substrings_in_text_with_spaces

logger = logging.getLogger(__name__)


def save_abstract_and_remaining_text(
    doc: TextDocumentWithLabeledPartitions, base_path: str
) -> None:
    abstract_annotations = [
        span for span in doc.labeled_partitions if span.label.lower().strip() == "abstract"
    ]
    if len(abstract_annotations) != 1:
        logger.warning(
            f"Expected exactly one abstract annotation, found {len(abstract_annotations)}"
        )
        return
    abstract_annotation = abstract_annotations[0]
    text_abstract = doc.text[abstract_annotation.start : abstract_annotation.end]
    text_remaining = doc.text[abstract_annotation.end :]
    with open(
        f"{base_path}.abstract.{abstract_annotation.start}_{abstract_annotation.end}.txt", "w"
    ) as f:
        f.write(text_abstract)
    with open(f"{base_path}.remaining.{abstract_annotation.end}.txt", "w") as f:
        f.write(text_remaining)


D_text = TypeVar("D_text", bound=TextBasedDocument)


def clean_doc(doc: D_text) -> D_text:
    # remove xml tags. Note that we also remove the Abstract tag, in contrast to the preprocessing
    # pipeline (see configs/dataset/sciarg_cleaned.yaml). This is because there, the abstracts are
    # removed at completely.
    doc = replace_substrings_in_text_with_spaces(
        doc,
        substrings=[
            "</H2>",
            "<H3>",
            "</Document>",
            "<H1>",
            "<H2>",
            "</H3>",
            "</H1>",
            "<Abstract>",
            "</Abstract>",
        ],
    )
    return doc


def main(out_dir: str, doc_id_whitelist: Optional[List[str]] = None) -> None:
    logger.info("Loading dataset from pie/sciarg")
    sciarg_with_abstracts = load_dataset(
        "pie/sciarg",
        revision="171478ce3c13cc484be5d7c9bc8f66d7d2f1c210",
        split="train",
    )
    if issubclass(sciarg_with_abstracts.document_type, BratDocument):
        ds_converted = sciarg_with_abstracts.to_document_type(
            TextDocumentWithLabeledMultiSpansBinaryRelationsAndLabeledPartitions
        )
    elif issubclass(sciarg_with_abstracts.document_type, BratDocumentWithMergedSpans):
        ds_converted = sciarg_with_abstracts.to_document_type(
            TextDocumentWithLabeledSpansBinaryRelationsAndLabeledPartitions
        )
    else:
        raise ValueError(f"Unsupported document type {sciarg_with_abstracts.document_type}")

    ds_clean = ds_converted.map(clean_doc)
    if doc_id_whitelist is not None:
        num_before = len(ds_clean)
        ds_clean = [doc for doc in ds_clean if doc.id in doc_id_whitelist]
        logger.info(
            f"Filtered dataset from {num_before} to {len(ds_clean)} documents based on doc_id_whitelist"
        )

    os.makedirs(out_dir, exist_ok=True)
    logger.info(f"Saving dataset to {out_dir}")
    for doc in ds_clean:
        save_abstract_and_remaining_text(doc, os.path.join(out_dir, doc.id))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Split SciArg dataset into abstract and remaining text"
    )
    parser.add_argument(
        "--out_dir",
        type=str,
        default="data/datasets/sciarg/abstracts_and_remaining_text",
        help="Path to save the split data",
    )
    parser.add_argument(
        "--doc_id_whitelist",
        type=str,
        nargs="+",
        default=["A32", "A33", "A34", "A35", "A36", "A37", "A38", "A39", "A40"],
        help="List of document ids to include in the split",
    )

    logging.basicConfig(level=logging.INFO)

    kwargs = vars(parser.parse_args())
    # allow for "all" to include all documents
    if len(kwargs["doc_id_whitelist"]) == 1 and kwargs["doc_id_whitelist"][0].lower() == "all":
        kwargs["doc_id_whitelist"] = None
    main(**kwargs)
    logger.info("Done")