"""BEIR loading from the Hugging Face Hub and population definition."""
from dataclasses import dataclass

from datasets import load_dataset


@dataclass
class Beir:
    name: str
    doc_ids: list      # corpus order
    doc_texts: list    # title + " " + text
    query_ids: list    # population order: test queries with >= 1 relevant doc
    query_texts: list
    qrels: list        # per population query: {doc_index: grade > 0}


def load(name):
    corpus = load_dataset(f"BeIR/{name}", "corpus", split="corpus")
    queries = load_dataset(f"BeIR/{name}", "queries", split="queries")
    qrels = load_dataset(f"BeIR/{name}-qrels", split="test")

    doc_ids = [str(x) for x in corpus["_id"]]
    doc_texts = [f"{t} {x}".strip() for t, x in zip(corpus["title"], corpus["text"])]
    doc_index = {d: i for i, d in enumerate(doc_ids)}

    rel = {}
    for q, d, s in zip(qrels["query-id"], qrels["corpus-id"], qrels["score"]):
        if s > 0 and str(d) in doc_index:
            rel.setdefault(str(q), {})[doc_index[str(d)]] = int(s)

    text = {str(q): t for q, t in zip(queries["_id"], queries["text"])}
    query_ids = sorted(q for q in rel if q in text)
    return Beir(name, doc_ids, doc_texts, query_ids, [text[q] for q in query_ids], [rel[q] for q in query_ids])