"""
Notes:

The dataset is quite small only containing 78 queries and 262 documents.
"""

from typing import Any

import datasets
from mteb.abstasks import AbsTaskRetrieval


class TwitterHjerneRetrieval(AbsTaskRetrieval):
    @property
    def description(self) -> dict[str, Any]:
        return {
            "name": "TwitterHjerneRetrieval",
            "hf_hub_name": "sorenmulli/da-hashtag-twitterhjerne",
            "description": "Danish question asked on Twitter with the Hashtag #Twitterhjerne ('Twitter brain') and their corresponding answer.",
            "reference": "https://huggingface.co/datasets/sorenmulli/da-hashtag-twitterhjerne",
            "type": "Retrieval",
            "category": "p2p",
            "eval_splits": ["train"],
            "eval_langs": ["da"],
            "main_score": "ndcg_at_10",
            "revision": "099ee143c7fdfa6bd7965be8c801cb161c313b29",
        }

    def load_data(self, **kwargs: dict):  # noqa: ARG002
        """
        Load dataset from HuggingFace hub
        """
        if self.data_loaded:
            return

        self.dataset: datasets.DatasetDict = datasets.load_dataset(
            self.description["hf_hub_name"],
            revision=self.description.get("revision"),
        )  # type: ignore

        self.dataset_transform()
        self.data_loaded = True

    def dataset_transform(self) -> None:
        """
        and transform to a retrieval datset, which have the following attributes

        self.corpus = Dict[doc_id, Dict[str, str]] #id => dict with document datas like title and text
        self.queries = Dict[query_id, str] #id => query
        self.relevant_docs = Dict[query_id, Dict[[doc_id, score]]
        """
        self.corpus = {}
        self.relevant_docs = {}
        self.queries = {}
        text2id = {}

        for split in self.dataset:
            ds: datasets.Dataset = self.dataset[split]  # type: ignore
            ds = ds.map(answers_to_list)

            self.queries[split] = {}
            self.relevant_docs[split] = {}
            self.corpus[split] = {}

            questions = ds["Question"]
            answers = ds["answers"]

            n = 0
            for q, answ in zip(questions, answers):
                if len(q.split(" ")) < 4 and answ:
                    continue
                query_id = str(n)
                self.queries[split][query_id] = q
                n += 1
                answer_ids = []
                for a in answ:
                    if a not in text2id:
                        text2id[a] = n
                        answer_id = str(n)
                        self.corpus[split][answer_id] = {"title": "", "text": a}
                        n += 1
                    else:
                        answer_id = str(text2id[a])
                    answer_ids.append(answer_id)

                self.relevant_docs[split][query_id] = {answer_id: 1 for answer_id in answer_ids}


def answers_to_list(example: dict) -> dict:
    example["answers"] = [v for k, v in example.items() if k.startswith("Answer") and v and len(v.split(" ")) > 3]
    return example
