"""Minimal label vocabulary: per-namespace string<->int maps, null label pinned to index 0.

Port of the one piece of AllenNLP's `Vocabulary` this project actually uses: separate NER /
relation label namespaces per `dataset` value (e.g. "radgraph-it__ner_labels"), with a
non-embedded null label "" always at index 0 so "argmax == 0" means "no label" throughout the
model. Built once from the training split and saved alongside the checkpoint so dev/test/
inference use the exact same label<->index mapping.
"""
import json
from pathlib import Path
from typing import Dict, Iterable, List

NULL_LABEL = ""


class Vocabulary:
    def __init__(self):
        self._token_to_index: Dict[str, Dict[str, int]] = {}
        self._index_to_token: Dict[str, Dict[int, str]] = {}

    def add_namespace(self, namespace: str) -> None:
        if namespace not in self._token_to_index:
            self._token_to_index[namespace] = {NULL_LABEL: 0}
            self._index_to_token[namespace] = {0: NULL_LABEL}

    def add_token(self, token: str, namespace: str) -> int:
        self.add_namespace(namespace)
        t2i = self._token_to_index[namespace]
        if token not in t2i:
            idx = len(t2i)
            t2i[token] = idx
            self._index_to_token[namespace][idx] = token
        return t2i[token]

    def get_token_index(self, token: str, namespace: str) -> int:
        return self._token_to_index[namespace][token]

    def get_token_from_index(self, index: int, namespace: str) -> str:
        return self._index_to_token[namespace][index]

    def get_vocab_size(self, namespace: str) -> int:
        return len(self._token_to_index[namespace])

    def get_namespaces(self) -> List[str]:
        return list(self._token_to_index.keys())

    @classmethod
    def from_documents(cls, documents: Iterable, ner_suffix="ner_labels",
                        relation_suffix="relation_labels") -> "Vocabulary":
        vocab = cls()
        for doc in documents:
            ner_ns = f"{doc.dataset}__{ner_suffix}"
            rel_ns = f"{doc.dataset}__{relation_suffix}"
            vocab.add_namespace(ner_ns)
            vocab.add_namespace(rel_ns)
            for sent in doc.sentences:
                for e in sent.ner or []:
                    vocab.add_token(e.label, ner_ns)
                for r in sent.relations or []:
                    vocab.add_token(r.label, rel_ns)
        return vocab

    def to_dict(self) -> dict:
        return {ns: dict(sorted(t2i.items(), key=lambda kv: kv[1]))
                for ns, t2i in self._token_to_index.items()}

    def save(self, path: str) -> None:
        Path(path).write_text(json.dumps(self.to_dict(), ensure_ascii=False, indent=2))

    @classmethod
    def load(cls, path: str) -> "Vocabulary":
        vocab = cls()
        for ns, t2i in json.loads(Path(path).read_text()).items():
            vocab.add_namespace(ns)
            for token, idx in sorted(t2i.items(), key=lambda kv: kv[1]):
                if token != NULL_LABEL:
                    vocab.add_token(token, ns)
        return vocab