from mteb.abstasks.classification import AbsTaskClassification
from mteb.abstasks.task_metadata import TaskMetadata

_LANGS = {
    "Danish": ["dan-Latn"],
    "Norwegian_b": ["nob-Latn"],
    "Norwegian_n": ["nno-Latn"],
    "Swedish": ["swe-Latn"],
}


class ScalaClassification(AbsTaskClassification):
    metadata = TaskMetadata(
        name="ScalaClassification",
        description="ScaLa a linguistic acceptability dataset for the mainland Scandinavian languages automatically constructed from dependency annotations in Universal Dependencies Treebanks. Published as part of 'ScandEval: A Benchmark for Scandinavian Natural Language Processing'",
        reference="https://aclanthology.org/2023.nodalida-1.20/",
        dataset={
            "path": "mteb/multilingual-scala-classification",
            "revision": "ec85bb6c69679ed15ac66c0bf6e180bf563eb137",
        },
        type="Classification",
        category="t2c",
        modalities=["text"],
        eval_splits=["test"],
        eval_langs=_LANGS,
        main_score="accuracy",
        date=(
            "1990-01-01",
            "2023-01-01",
        ),  # derived from dependency treebank, this a the best guess
        domains=["Fiction", "News", "Non-fiction", "Blog", "Spoken", "Web", "Written"],
        task_subtypes=["Linguistic acceptability"],
        license="cc-by-sa-4.0",
        annotations_creators="human-annotated",
        dialect=[],
        sample_creation="created",
        bibtex_citation=r"""
@inproceedings{nielsen-2023-scandeval,
  address = {T{\'o}rshavn, Faroe Islands},
  author = {Nielsen, Dan},
  booktitle = {Proceedings of the 24th Nordic Conference on Computational Linguistics (NoDaLiDa)},
  editor = {Alum{\"a}e, Tanel  and
Fishel, Mark},
  month = may,
  pages = {185--201},
  publisher = {University of Tartu Library},
  title = {{S}cand{E}val: A Benchmark for {S}candinavian Natural Language Processing},
  url = {https://aclanthology.org/2023.nodalida-1.20},
  year = {2023},
}
""",
        prompt="Classify passages in Scandinavian Languages based on linguistic acceptability",
    )

    samples_per_label = 32

    def dataset_transform(
        self,
        num_proc: int | None = None,
    ):
        for lang in self.dataset.keys():
            # convert label to a 0/1 label
            labels = self.dataset[lang]["train"]["label"]
            lab2idx = {lab: idx for idx, lab in enumerate(set(labels))}
            self.dataset[lang] = self.dataset[lang].map(
                lambda x: {"label": lab2idx[x["label"]]}, remove_columns=["label"]
            )
