from __future__ import annotations

from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Any, List, Tuple

if TYPE_CHECKING:
    from data.kg_dataset import KGDataset


class BaseMetric(ABC):
    """Base class for metrics that compute Weisfeiler-Lehman Entropy
        Complexity (WLEC) or similar metrics.
    """

    def __init__(self, dataset: KGDataset) -> None:
        self.dataset = dataset
        self.num_entities = dataset.num_entities

        self.out_adj, self.in_adj = self._build_adjacency_lists()

    def _build_adjacency_lists(
        self,
    ) -> Tuple[List[List[Tuple[int, int]]], List[List[Tuple[int, int]]]]:
        """Create *in* and *out* adjacency lists from mapped triples.

        Parameters
        ----------
        triples:
            Tensor of shape *(m, 3)* containing *(h, r, t)* triples (``dtype=torch.long``).
        num_entities:
            Total number of entities.  Needed to allocate adjacency containers.

        Returns
        -------
        out_adj, in_adj
            Each element ``out_adj[v]`` is a list ``[(r, t), ...]``.  Each element
            ``in_adj[v]`` is a list ``[(r, h), ...]``.
        """

        triples = self.dataset.triples  # (m, 3)
        num_entities = self.num_entities

        out_adj: List[List[Tuple[int, int]]] = [[] for _ in range(num_entities)]
        in_adj: List[List[Tuple[int, int]]] = [[] for _ in range(num_entities)]
        for h, r, t in triples.tolist():
            out_adj[h].append((r, t))
            in_adj[t].append((r, h))
        return out_adj, in_adj

    @abstractmethod
    def compute(
        self,
        *args,
        **kwargs,
    ) -> Any: ...