src/gcmd_classifier/vocabulary/index.py

"""In-memory canonical vocabulary index for GCMD Science Keywords."""

from __future__ import annotations

from collections import Counter
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType

from gcmd_classifier.errors import VocabularyLookupError
from gcmd_classifier.models import CanonicalConceptRecord, HierarchyLevel


@dataclass(frozen=True)
class VocabularyIndex:
    """Lookup structures derived from the canonical GCMD hierarchy."""

    records_by_uuid: Mapping[str, CanonicalConceptRecord]
    records_by_path: Mapping[str, CanonicalConceptRecord]
    parent_by_uuid: Mapping[str, str | None]
    children_by_uuid: Mapping[str | None, tuple[str, ...]]
    topic_uuids: tuple[str, ...]
    terms_by_topic_uuid: Mapping[str, tuple[str, ...]]
    variables_by_parent_uuid: Mapping[str, tuple[str, ...]]
    vocabulary_version: str
    root_level: str
    root_name: str

    def __post_init__(self) -> None:
        object.__setattr__(self, "records_by_uuid", MappingProxyType(dict(self.records_by_uuid)))
        object.__setattr__(self, "records_by_path", MappingProxyType(dict(self.records_by_path)))
        object.__setattr__(self, "parent_by_uuid", MappingProxyType(dict(self.parent_by_uuid)))
        object.__setattr__(self, "children_by_uuid", MappingProxyType(dict(self.children_by_uuid)))
        object.__setattr__(
            self, "terms_by_topic_uuid", MappingProxyType(dict(self.terms_by_topic_uuid))
        )
        object.__setattr__(
            self, "variables_by_parent_uuid", MappingProxyType(dict(self.variables_by_parent_uuid))
        )

    @property
    def records(self) -> tuple[CanonicalConceptRecord, ...]:
        """All UUID-bearing canonical records in deterministic path order."""
        return tuple(self.records_by_path[path] for path in sorted(self.records_by_path))

    def __len__(self) -> int:
        return len(self.records_by_uuid)

    def get(self, uuid: str) -> CanonicalConceptRecord:
        """Return a concept by UUID or raise a typed lookup error."""
        try:
            return self.records_by_uuid[uuid]
        except KeyError as exc:
            raise VocabularyLookupError(f"Unknown GCMD UUID: {uuid}") from exc

    def get_by_path(self, canonical_path: str) -> CanonicalConceptRecord:
        """Return a concept by canonical path or raise a typed lookup error."""
        try:
            return self.records_by_path[canonical_path]
        except KeyError as exc:
            raise VocabularyLookupError(f"Unknown GCMD canonical path: {canonical_path}") from exc

    def parent_of(self, uuid: str) -> str | None:
        """Return the parent UUID for a concept, or None for Topic records."""
        self.get(uuid)
        return self.parent_by_uuid[uuid]

    def children_of(self, parent_uuid: str | None) -> tuple[str, ...]:
        """Return direct child UUIDs for a parent UUID; None returns Topic UUIDs."""
        if parent_uuid is not None:
            self.get(parent_uuid)
        return self.children_by_uuid.get(parent_uuid, ())

    def topics(self) -> tuple[CanonicalConceptRecord, ...]:
        """Return all Topic records in source order."""
        return tuple(self.get(uuid) for uuid in self.topic_uuids)

    def terms_for_topic(self, topic_uuid: str) -> tuple[CanonicalConceptRecord, ...]:
        """Return direct Term records under a Topic UUID."""
        topic = self.get(topic_uuid)
        if topic.level != "Topic":
            raise VocabularyLookupError(f"UUID is not a Topic: {topic_uuid}")
        return tuple(self.get(uuid) for uuid in self.terms_by_topic_uuid.get(topic_uuid, ()))

    def variables_for_parent(self, parent_uuid: str) -> tuple[CanonicalConceptRecord, ...]:
        """Return direct Variable records under a Term or Variable parent."""
        parent = self.get(parent_uuid)
        if parent.level == "Topic":
            raise VocabularyLookupError(
                f"Topic UUID does not have variable children: {parent_uuid}"
            )
        return tuple(self.get(uuid) for uuid in self.variables_by_parent_uuid.get(parent_uuid, ()))

    def ancestors_of(self, uuid: str) -> tuple[str, ...]:
        """Return ancestor UUIDs from nearest parent to highest UUID-bearing ancestor."""
        self.get(uuid)
        ancestors: list[str] = []
        current = self.parent_by_uuid[uuid]
        while current is not None:
            ancestors.append(current)
            current = self.parent_by_uuid[current]
        return tuple(ancestors)

    def is_ancestor(self, ancestor_uuid: str, descendant_uuid: str) -> bool:
        """Return whether the first UUID is an ancestor of the second UUID."""
        self.get(ancestor_uuid)
        return ancestor_uuid in self.ancestors_of(descendant_uuid)

    def is_descendant(self, descendant_uuid: str, ancestor_uuid: str) -> bool:
        """Return whether the first UUID is a descendant of the second UUID."""
        return self.is_ancestor(ancestor_uuid, descendant_uuid)

    def count_by_level(self) -> Counter[HierarchyLevel]:
        """Return UUID-bearing record counts by hierarchy level."""
        return Counter(record.level for record in self.records_by_uuid.values())