<filename>src/ilp/data/statement_factory.py
# -*- coding: utf-8 -*-

"""A graph factory that creates the PyTorchGeometric data object, and holds references to the textual descriptions."""
import functools
import logging
import pathlib
import re
from collections import defaultdict
from typing import Collection, Dict, Iterable, List, Mapping, Optional, Tuple, Union

import numpy as np
import torch
from torch_geometric.data import Data
from tqdm import tqdm

from .features import dataset_resolver, get_node_features, get_remapped_node_features
from .instances import LCWAInstances, SLCWAInstances
from .split import filter_qualifiers
from ..utils import PADDING, load_rdf_star_statements

logger = logging.getLogger(__name__)

INVERSE_SUFFIX = '_inverse'
TRIPLES_DF_COLUMNS = ('head_id', 'head_label', 'relation_id', 'relation_label', 'tail_id', 'tail_label')


def _map_statements(
    statements: np.ndarray,
    entity_to_id,
    relation_to_id,
    max_num_qualifier_pairs: int,
) -> Tuple[torch.LongTensor, Collection[np.ndarray]]:
    values = set(list(entity_to_id.keys()) + list(relation_to_id.keys()))
    statements_filtered = []
    mapped_statements = []
    max_length = max_num_qualifier_pairs * 2 + 3
    for s in statements:
        unique_values_in_s = set(s).difference(values)
        if len(unique_values_in_s) != 0:
            continue
        statements_filtered.append(s)
        # Map to id
        mapped_s = []
        for i, v in enumerate(s):
            if i % 2 == 0:
                mapped_s.append(entity_to_id[v])
            else:
                mapped_s.append(relation_to_id[v])
        # Pad statement
        if max_length >= len(mapped_s):
            stop = max_length - len(mapped_s)
            mapped_s = mapped_s + [entity_to_id[PADDING] for _ in range(stop)]
        else:
            mapped_s = mapped_s[:max_length]

        mapped_statements.append(mapped_s)
    mapped_statements = torch.as_tensor(data=mapped_statements, dtype=torch.long)
    return mapped_statements, statements_filtered


class StatementsFactory:
    """Create graph and save it as torch_geometric.data."""

    mapped_qualifiers: torch.LongTensor
    triple_to_idx: Mapping[Tuple[int, int, int], Collection[int]]
    head_relation_to_id: Mapping[Tuple[int, int], Collection[int]]
    tail_relation_to_id: Mapping[Tuple[int, int], Collection[int]]
    padding_idx = 0

    def __init__(
        self,
        path: Union[None, str, pathlib.Path, List[str], List[pathlib.Path]] = None,
        dataset_name: str = None,
        statements: List[List[str]] = None,
        max_num_qualifier_pairs: int = None,
        create_inverse_triples: bool = False,
        entity_to_id: Mapping[str, int] = None,
        relation_to_id: Mapping[str, int] = None,
        use_node_features: bool = True,
        feature_dim: Optional[int] = None,
        enforce_node_features: bool = False,
        is_inductive: bool = False,
    ) -> None:
        """
        Initialize the factory.

        :param path:
            path to the graph file
        :param dataset_name:
            The dataset name.
        :param statements:
            directly input triples (optional)
        :param entity_to_id:
            mapping from labels to a numeric id. If None, infer from statements.
        :param relation_to_id:
            mapping relation type to an numeric id. If None, infer from statements.
        :param use_node_features:
            whether to use node features or generate random ones.
        :param feature_dim:
            dimension of random features for nodes and edges.
        """
        self.path = path
        if path is not None:
            self.statements = load_rdf_star_statements(path=path, max_len=max_num_qualifier_pairs * 2 + 3)
        else:
            if statements is None:
                raise ValueError("You must explictly provide statements if no path is given.")
            self.statements = statements

        self.create_inverse_triples = create_inverse_triples
        self._process_inverse_relations()
        self.max_num_qualifier_pairs = max_num_qualifier_pairs

        if use_node_features:
            drive_label_to_id = None
            emb = None
            if enforce_node_features:
                drive_label_to_id, emb = get_node_features(
                    dataset=dataset_name,
                    cache_root=None,
                    force=False,
                )

                nodes_with_features = set(drive_label_to_id.keys())

                # Filter statements based on node features
                filtered_statements = []

                known_relations = None if relation_to_id is None else set(relation_to_id.keys())

                for s in self.statements:
                    # convert to tuple
                    s = tuple(s)
                    # if head or tail have no features -> skip statement
                    if not nodes_with_features.issuperset(s[:3:2]):
                        continue
                    # if any relation is unknown:
                    if is_inductive:
                        if s[1] not in known_relations:
                            continue
                    # filter qualifiers by entities
                    hrt, qs = s[:3], s[3:]
                    qs = filter_qualifiers(qualifiers=qs, allowed=nodes_with_features, entities=True)
                    # filter qualifiers by relations
                    if known_relations is not None:
                        qs = filter_qualifiers(qualifiers=qs, allowed=known_relations, entities=False)
                    filtered_statements.append(hrt + qs)

                self.statements = filtered_statements

        # Generate entity mapping if necessary
        if not is_inductive:
            if entity_to_id is None or relation_to_id is None:
                self._create_mappings_rdf_star()
            else:
                self.entity_to_id = entity_to_id
                self.relation_to_id = relation_to_id
        else:
            if entity_to_id is None:
                self._create_mappings_rdf_star(create_relation_to_id=False)
            else:
                self.entity_to_id = entity_to_id
            if relation_to_id is None:
                raise ValueError("A inductive setting is defined, but the relation_to_id mapping is not provided.")

            self.relation_to_id = relation_to_id

        # Note: Mapped statements are not padded
        self.mapped_statements, self.statements_filtered = _map_statements(
            statements=self.statements,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            max_num_qualifier_pairs=max_num_qualifier_pairs,
        )

        # load node features
        if use_node_features:
            self.node_feature_tensor = torch.as_tensor(
                data=get_remapped_node_features(
                    dataset=dataset_name,
                    label_to_id=self.entity_to_id,
                    drive_label_to_id=drive_label_to_id,
                    emb=emb,
                    # TODO: Allow to provide this via constructor parameter?
                    cache_root=None,
                    force=False,
                ),
                dtype=torch.get_default_dtype(),
            )
            dim = self.node_feature_tensor.shape[-1]
            if feature_dim is not None and dim != feature_dim:
                logger.warning(
                    f"Using pre-trained features with dim {dim}, although feature_dim={feature_dim} was passed. "
                    f"To disable this warning pass feature_dim=None."
                )
        else:
            if feature_dim is None:
                raise ValueError("Must provide feature_dim if random features shall be used.")
            self.node_feature_tensor = torch.rand(
                self.num_entities,
                feature_dim,
            )
        self.feature_dim = self.node_feature_tensor.shape[-1]
        self.max_num_qualifier_pairs = max_num_qualifier_pairs

        # store non qualifier only entities
        self.non_qualifier_only_entities = self.mapped_statements[:, [0, 2]].unique()

    def _process_inverse_relations(self):
        relations = list(sorted(set([e for s in self.statements for e in s[1::2]])))

        # Check if the triples are inverted already
        relations_already_inverted = self._check_already_inverted_relations(relations)

        if self.create_inverse_triples or relations_already_inverted:
            self.create_inverse_triples = True
            if relations_already_inverted:
                logger.info(
                    f'Some triples already have suffix {INVERSE_SUFFIX}. '
                    f'Creating TriplesFactory based on inverse triples',
                )
                self.relation_to_inverse = {
                    re.sub('_inverse$', '', relation): f"{re.sub('_inverse$', '', relation)}{INVERSE_SUFFIX}"
                    for relation in relations
                }

            else:
                self.relation_to_inverse = {
                    relation: f"{relation}{INVERSE_SUFFIX}"
                    for relation in relations
                }
                inverse_statements = [[s[2], self.relation_to_inverse[s[1]], s[0], *s[3:]] for s in self.statements]
                # extend original triples with inverse ones
                self.statements.extend(inverse_statements)

        else:
            self.create_inverse_triples = False
            self.relation_to_inverse = None

    @staticmethod
    def _check_already_inverted_relations(relations: Iterable[str]) -> bool:
        for relation in relations:
            if relation.endswith(INVERSE_SUFFIX):
                # We can terminate the search after finding the first inverse occurrence
                return True

        return False

    def create_data_object(self) -> Data:
        """Create PyTorch Geometric data object."""
        node_features = self._create_node_feature_tensor()
        edge_index = self._create_edge_index()
        qualifier_index = self._create_qualifier_index()
        edge_type = self._create_edge_type()

        entities = torch.as_tensor(
            data=sorted(set(self.entity_to_id.values())),  # exclude padding entity
            dtype=torch.long,
        )

        return Data(
            x=node_features,
            edge_index=edge_index,
            qualifier_index=qualifier_index,
            edge_type=edge_type,
            entities=entities,
        )

    def _compose_qualifier_batch(self, keys, mapping) -> torch.LongTensor:
        # batch.shape:
        batch_size = len(keys)

        # Lookup IDs of qualifiers
        qualifier_ids = [
            mapping.get(tuple(k.tolist()))
            for k in keys
        ]

        # Determine maximum length (for dynamic padding)
        max_len = max(map(len, qualifier_ids))

        # bound maximum length for guaranteed max memory consumption
        # TODO: num_qualifier_pairs = max_num_qualifier_pairs ?
        max_len = min(max_len, self.max_num_qualifier_pairs)

        # Allocate result
        result = torch.full(size=(batch_size, max_len), fill_value=-1, dtype=torch.long)

        # Retrieve qualifiers
        for i, this_qualifier_ids in enumerate(qualifier_ids):
            # limit number of qualifiers
            # TODO: shuffle?
            this_qualifier_ids = this_qualifier_ids[:max_len]
            result[i, :len(this_qualifier_ids)] = torch.as_tensor(data=this_qualifier_ids, dtype=torch.long)

        return result

    def _create_qualifier_index(self):
        """
        Create a COO matrix of shape (3, num_qualifiers). Only non-zero (non-padded) qualifiers are retained.
        row0: qualifying relations
        row1: qualifying entities
        row2: index row which connects a pair (qual_r, qual_e) to a statement index k
        :return: shape: (3, num_qualifiers)
        """
        if self.max_num_qualifier_pairs is None or self.max_num_qualifier_pairs == 0:
            return None

        qual_relations, qual_entities, qual_k = [], [], []

        # It is assumed that statements are already padded
        for triple_id, statement in enumerate(self.mapped_statements):
            qualifiers = statement[3:]
            entities = qualifiers[1::2]
            relations = qualifiers[::2]
            # Ensure that PADDING has id=0
            non_zero_rels = relations[np.nonzero(relations)]
            non_zero_ents = entities[np.nonzero(entities)]
            assert len(non_zero_rels) == len(non_zero_ents), \
                "Number of non-padded qualifying relations is not equal to the # of qualifying entities"
            num_qualifier_pairs = non_zero_ents.shape[0]

            for j in range(num_qualifier_pairs):
                qual_relations.append(non_zero_rels[j].item())
                qual_entities.append(non_zero_ents[j].item())
                qual_k.append(triple_id)

        qualifier_index = torch.stack([
            torch.tensor(qual_relations, dtype=torch.long),
            torch.tensor(qual_entities, dtype=torch.long),
            torch.tensor(qual_k, dtype=torch.long)
        ], dim=0)

        if self.create_inverse_triples:
            # qualifier index is the same for inverse statements
            qualifier_index[2, len(qual_relations) // 2:] = qualifier_index[2, :len(qual_relations) // 2]

        return qualifier_index

    def _create_node_feature_tensor(self) -> torch.Tensor:
        """Create the node feature tensor."""
        if self.node_feature_tensor is not None:
            return self.node_feature_tensor
        if self.one_hot_encoding:
            return torch.eye(n=len(self.entity_to_id))
        else:
            return torch.empty(len(self.entity_to_id), self.feature_dim)

    def _create_edge_index(self) -> torch.Tensor:
        """Create edge index where first row represents the source nodes and the second row the target nodes."""
        mapped_heads = self.mapped_statements[:, 0].view(1, -1)
        mapped_tails = self.mapped_statements[:, 2].view(1, -1)
        edge_index = torch.cat([mapped_heads, mapped_tails], dim=0)
        return edge_index

    def _create_edge_type(self) -> torch.Tensor:
        """Create edge type tensor where each entry correspond to the relationship type of a triple in the dataset."""
        # Inverse triples are created in the base class

        return self.mapped_statements[:, 1]

    def _create_mappings_rdf_star(
        self,
        create_entity_to_id: bool = True,
        create_relation_to_id: bool = True,
    ):
        """."""

        max_length = self.max_num_qualifier_pairs * 2 + 3
        # Ensure that no entity/relation is part of the mappings that will not be seen
        statements = [s[:max_length] if len(s) > max_length else s for s in self.statements]

        if create_relation_to_id:
            relations = [PADDING] + list(sorted(set([r for s in statements for r in s[1::2]])))
            if self.create_inverse_triples:
                relations = [PADDING] + [elem for pair in
                                         zip(list(self.relation_to_inverse.keys()),
                                             list(self.relation_to_inverse.values()))
                                         for elem in pair]
                # old
                # relations = [PADDING] + list(self.relation_to_inverse.keys()) + (list(self.relation_to_inverse.values()))
            # When creating inverse relations, we will also generate padding_inverse which will never be used
            # unless we use score_r()
            self.relation_to_id = {relation: id for id, relation in enumerate(relations)}

        if create_entity_to_id:
            entities = [PADDING] + list(sorted(set([e for s in statements for e in s[::2]])))
            self.entity_to_id = {entity: id for id, entity in enumerate(entities)}

    def create_lcwa_instances(self, use_tqdm: Optional[bool] = None) -> LCWAInstances:
        """Create LCWA instances for this factory's statements."""
        s_p_q_to_multi_tails = _create_multi_label_tails_instance(
            mapped_statements=self.mapped_statements,
            use_tqdm=use_tqdm,
        )
        spq, multi_o = zip(*s_p_q_to_multi_tails.items())
        mapped_statements: torch.LongTensor = torch.tensor(spq, dtype=torch.long)
        labels = np.array([np.array(item) for item in multi_o], dtype=object)

        # create mask
        entity_mask = torch.zeros(self.num_entities, dtype=torch.bool)
        entity_mask[self.non_qualifier_only_entities] = True

        return LCWAInstances(
            mapped_statements=mapped_statements,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            labels=labels,
            # entity_mask=entity_mask,
        )

    def create_slcwa_instances(self) -> SLCWAInstances:
        """Create sLCWA instances for this factory's statements."""
        return SLCWAInstances(
            mapped_statements=self.mapped_statements,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            entities=self.non_qualifier_only_entities,
        )

    @property
    def num_entities(self) -> int:
        """The number of unique entities."""
        return len(self.entity_to_id)

    @property
    def num_relations(self) -> int:
        """The number of unique relations."""
        return len(self.relation_to_id)

    @property
    def statement_length(self) -> int:
        """The number of unique relations."""
        return self.max_num_qualifier_pairs * 2 + 3

    @property
    def num_statements(self) -> int:
        """The number of statements."""
        return self.mapped_statements.shape[0]

    @property
    def qualifier_ratio(self) -> float:
        """Return the percentage of statements with qualifiers."""
        return (~(self.mapped_statements[:, 3::2] == self.padding_idx).all(dim=1)).sum().item() / self.num_statements

    def extra_repr(self) -> str:
        """Extra representation."""
        return f"num_entities={self.num_entities:,}, " \
               f"num_relations={self.num_relations:,}, " \
               f"num_statements={self.num_statements:,}, " \
               f"max_num_qualifier_pairs={self.max_num_qualifier_pairs}, " \
               f"qualifier_ratio={self.qualifier_ratio:.2%}"

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self.extra_repr()})"

    # compatability with pykeen
    @property
    def num_triples(self) -> int:
        return self.num_statements


def _create_multi_label_tails_instance(
    mapped_statements: torch.Tensor,
    use_tqdm: Optional[bool] = None,
) -> Dict[Tuple[int, int], List[int]]:
    """Create for each (h,r,q*) pair the multi tail label."""
    logger.debug('Creating multi label tails instance')

    '''
    The mapped triples matrix has to be a numpy array to ensure correct pair hashing, as explained in
    https://github.com/pykeen/pykeen/commit/1bc71fe4eb2f24190425b0a4d0b9d6c7b9c4653a
    '''
    mapped_statements = mapped_statements.cpu().detach().numpy()

    s_p_q_to_multi_tails_new = _create_multi_label_instances(
        mapped_statements,
        element_1_index=0,
        element_2_index=1,
        label_index=2,
        use_tqdm=use_tqdm,
    )

    logger.debug('Created multi label tails instance')

    return s_p_q_to_multi_tails_new


def _create_multi_label_instances(
    mapped_statements: torch.Tensor,
    element_1_index: int,
    element_2_index: int,
    label_index: int,
    use_tqdm: Optional[bool] = None,
) -> Dict[Tuple[int, ...], List[int]]:
    """Create for each (element_1, element_2) pair the multi-label."""
    instance_to_multi_label = defaultdict(set)

    if use_tqdm is None:
        use_tqdm = True

    it = mapped_statements

    if use_tqdm:
        it = tqdm(mapped_statements, unit='statement', unit_scale=True, desc='Grouping statements')

    for row in it:
        instance = tuple([row[element_1_index], row[element_2_index]] + row[3:].tolist())
        instance_to_multi_label[instance].add(
            row[label_index])

    # Create lists out of sets for proper numpy indexing when loading the labels
    # TODO is there a need to have a canonical sort order here?
    instance_to_multi_label_new = {
        key: list(value)
        for key, value in instance_to_multi_label.items()
    }

    return instance_to_multi_label_new


def get_semi_inductive_factories(
    dataset_name: str,
    max_num_qualifier_pairs: Optional[int] = None,
    create_inverse_triples: bool = True,
) -> Tuple[StatementsFactory, StatementsFactory, StatementsFactory]:
    """Get factories for semi-inductive split."""
    dataset = dataset_resolver.make(query=dataset_name)
    paths = {
        key: dataset.get_path(split="semi_inductive", part=key)
        for key in [
            "train",
            "validation",
            "test",
        ]
    }

    # Step 1: Get relation_to_id based on relations in train
    tmp_stmt_factory = StatementsFactory(
        path=paths['train'],
        dataset_name=dataset_name,
        use_node_features=True,
        is_inductive=False,
        enforce_node_features=True,
        create_inverse_triples=create_inverse_triples,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )

    relation_to_id = tmp_stmt_factory.relation_to_id

    # Step 2: Create entity_to_id based on all entities. Entities need to be connected to train-relations
    tmp_stmt_factory = StatementsFactory(
        path=list(paths.values()),
        dataset_name=dataset_name,
        create_inverse_triples=create_inverse_triples,
        relation_to_id=relation_to_id,
        use_node_features=True,
        is_inductive=True,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Full: {tmp_stmt_factory}")

    entity_to_id = tmp_stmt_factory.entity_to_id

    del tmp_stmt_factory

    # Create statement factories
    train_sf = StatementsFactory(
        path=paths["train"],
        dataset_name=dataset_name,
        entity_to_id=entity_to_id,
        relation_to_id=relation_to_id,
        create_inverse_triples=create_inverse_triples,
        use_node_features=True,
        is_inductive=False,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Training: {train_sf}")

    validation_sf = StatementsFactory(
        path=paths["validation"],
        dataset_name=dataset_name,
        is_inductive=True,
        create_inverse_triples=False,
        # inductive entities
        entity_to_id=train_sf.entity_to_id,
        # shared relations
        relation_to_id=train_sf.relation_to_id,
        use_node_features=True,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Validation: {validation_sf}")

    # In fully inductive setting, validation graph consists of training entities
    test_sf = StatementsFactory(
        path=paths["test"],
        dataset_name=dataset_name,
        create_inverse_triples=False,
        is_inductive=True,
        # inductive entities
        entity_to_id=train_sf.entity_to_id,
        # inductive relations
        relation_to_id=train_sf.relation_to_id,
        use_node_features=True,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Testing: {test_sf}")

    return train_sf, validation_sf, test_sf


def get_fully_inductive_factories(
    dataset_name: str,
    dataset_version: str,
    max_num_qualifier_pairs: int,
    create_inverse_triples: bool,
) -> Tuple[StatementsFactory, StatementsFactory, StatementsFactory, StatementsFactory]:
    """Create statement factories for the fully inductive setting."""
    dataset = dataset_resolver.make(query=dataset_name)
    paths = {
        key: dataset.get_path(split="inductive", part=key, version=dataset_version)
        for key in [
            "train",
            "inference",
            "validation",
            "test",
        ]
    }
    train_transductive_sf = StatementsFactory(
        path=paths["train"],
        dataset_name=dataset_name,
        create_inverse_triples=create_inverse_triples,
        use_node_features=True,
        is_inductive=False,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    train_inductive_sf = StatementsFactory(
        path=paths["inference"],
        dataset_name=dataset_name,
        create_inverse_triples=create_inverse_triples,
        relation_to_id=train_transductive_sf.relation_to_id,
        is_inductive=True,
        use_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    # In fully inductive setting, validation graph consists of training entities
    validation_sf = StatementsFactory(
        path=paths["validation"],
        dataset_name=dataset_name,
        entity_to_id=train_inductive_sf.entity_to_id,
        relation_to_id=train_inductive_sf.relation_to_id,
        use_node_features=True,
        is_inductive=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    # Test graph contains only unseen entities, therefore, create
    test_sf = StatementsFactory(
        path=paths["test"],
        dataset_name=dataset_name,
        entity_to_id=train_inductive_sf.entity_to_id,
        relation_to_id=train_inductive_sf.relation_to_id,
        use_node_features=True,
        is_inductive=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    return train_transductive_sf, train_inductive_sf, test_sf, validation_sf


# buffer loading (for HPO)
@functools.lru_cache(maxsize=1)
def get_factories(
    dataset_name: str,
    dataset_version: Optional[str],
    inductive_setting: str,
    max_num_qualifier_pairs: int,
    create_inverse_triples: bool,
) -> Tuple[StatementsFactory, StatementsFactory, StatementsFactory, StatementsFactory]:
    """Load triples factories."""
    if inductive_setting == "semi":
        train_sf, validation_sf, test_sf = get_semi_inductive_factories(
            dataset_name=dataset_name,
            max_num_qualifier_pairs=max_num_qualifier_pairs,
            create_inverse_triples=create_inverse_triples,
        )
        inference_sf = train_sf

    elif inductive_setting == "full":
        train_sf, inference_sf, test_sf, validation_sf = get_fully_inductive_factories(
            dataset_name=dataset_name,
            dataset_version=dataset_version,
            max_num_qualifier_pairs=max_num_qualifier_pairs,
            create_inverse_triples=create_inverse_triples,
        )

    else:
        raise ValueError(f"Unknown inductive setting: {inductive_setting}")

    return inference_sf, test_sf, train_sf, validation_sf

<filename>src/ilp/data/statement_factory.py
# -*- coding: utf-8 -*-

"""A graph factory that creates the PyTorchGeometric data object, and holds references to the textual descriptions."""
import functools
import logging
import pathlib
import re
from collections import defaultdict
from typing import Collection, Dict, Iterable, List, Mapping, Optional, Tuple, Union

import numpy as np
import torch
from torch_geometric.data import Data
from tqdm import tqdm

from .features import dataset_resolver, get_node_features, get_remapped_node_features
from .instances import LCWAInstances, SLCWAInstances
from .split import filter_qualifiers
from ..utils import PADDING, load_rdf_star_statements

logger = logging.getLogger(__name__)

INVERSE_SUFFIX = '_inverse'
TRIPLES_DF_COLUMNS = ('head_id', 'head_label', 'relation_id', 'relation_label', 'tail_id', 'tail_label')


def _map_statements(
    statements: np.ndarray,
    entity_to_id,
    relation_to_id,
    max_num_qualifier_pairs: int,
) -> Tuple[torch.LongTensor, Collection[np.ndarray]]:
    values = set(list(entity_to_id.keys()) + list(relation_to_id.keys()))
    statements_filtered = []
    mapped_statements = []
    max_length = max_num_qualifier_pairs * 2 + 3
    for s in statements:
        unique_values_in_s = set(s).difference(values)
        if len(unique_values_in_s) != 0:
            continue
        statements_filtered.append(s)
        # Map to id
        mapped_s = []
        for i, v in enumerate(s):
            if i % 2 == 0:
                mapped_s.append(entity_to_id[v])
            else:
                mapped_s.append(relation_to_id[v])
        # Pad statement
        if max_length >= len(mapped_s):
            stop = max_length - len(mapped_s)
            mapped_s = mapped_s + [entity_to_id[PADDING] for _ in range(stop)]
        else:
            mapped_s = mapped_s[:max_length]

        mapped_statements.append(mapped_s)
    mapped_statements = torch.as_tensor(data=mapped_statements, dtype=torch.long)
    return mapped_statements, statements_filtered


class StatementsFactory:
    """Create graph and save it as torch_geometric.data."""

    mapped_qualifiers: torch.LongTensor
    triple_to_idx: Mapping[Tuple[int, int, int], Collection[int]]
    head_relation_to_id: Mapping[Tuple[int, int], Collection[int]]
    tail_relation_to_id: Mapping[Tuple[int, int], Collection[int]]
    padding_idx = 0

    def __init__(
        self,
        path: Union[None, str, pathlib.Path, List[str], List[pathlib.Path]] = None,
        dataset_name: str = None,
        statements: List[List[str]] = None,
        max_num_qualifier_pairs: int = None,
        create_inverse_triples: bool = False,
        entity_to_id: Mapping[str, int] = None,
        relation_to_id: Mapping[str, int] = None,
        use_node_features: bool = True,
        feature_dim: Optional[int] = None,
        enforce_node_features: bool = False,
        is_inductive: bool = False,
    ) -> None:
        """
        Initialize the factory.

        :param path:
            path to the graph file
        :param dataset_name:
            The dataset name.
        :param statements:
            directly input triples (optional)
        :param entity_to_id:
            mapping from labels to a numeric id. If None, infer from statements.
        :param relation_to_id:
            mapping relation type to an numeric id. If None, infer from statements.
        :param use_node_features:
            whether to use node features or generate random ones.
        :param feature_dim:
            dimension of random features for nodes and edges.
        """
        self.path = path
        if path is not None:
            self.statements = load_rdf_star_statements(path=path, max_len=max_num_qualifier_pairs * 2 + 3)
        else:
            if statements is None:
                raise ValueError("You must explictly provide statements if no path is given.")
            self.statements = statements

        self.create_inverse_triples = create_inverse_triples
        self._process_inverse_relations()
        self.max_num_qualifier_pairs = max_num_qualifier_pairs

        if use_node_features:
            drive_label_to_id = None
            emb = None
            if enforce_node_features:
                drive_label_to_id, emb = get_node_features(
                    dataset=dataset_name,
                    cache_root=None,
                    force=False,
                )

                nodes_with_features = set(drive_label_to_id.keys())

                # Filter statements based on node features
                filtered_statements = []

                known_relations = None if relation_to_id is None else set(relation_to_id.keys())

                for s in self.statements:
                    # convert to tuple
                    s = tuple(s)
                    # if head or tail have no features -> skip statement
                    if not nodes_with_features.issuperset(s[:3:2]):
                        continue
                    # if any relation is unknown:
                    if is_inductive:
                        if s[1] not in known_relations:
                            continue
                    # filter qualifiers by entities
                    hrt, qs = s[:3], s[3:]
                    qs = filter_qualifiers(qualifiers=qs, allowed=nodes_with_features, entities=True)
                    # filter qualifiers by relations
                    if known_relations is not None:
                        qs = filter_qualifiers(qualifiers=qs, allowed=known_relations, entities=False)
                    filtered_statements.append(hrt + qs)

                self.statements = filtered_statements

        # Generate entity mapping if necessary
        if not is_inductive:
            if entity_to_id is None or relation_to_id is None:
                self._create_mappings_rdf_star()
            else:
                self.entity_to_id = entity_to_id
                self.relation_to_id = relation_to_id
        else:
            if entity_to_id is None:
                self._create_mappings_rdf_star(create_relation_to_id=False)
            else:
                self.entity_to_id = entity_to_id
            if relation_to_id is None:
                raise ValueError("A inductive setting is defined, but the relation_to_id mapping is not provided.")

            self.relation_to_id = relation_to_id

        # Note: Mapped statements are not padded
        self.mapped_statements, self.statements_filtered = _map_statements(
            statements=self.statements,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            max_num_qualifier_pairs=max_num_qualifier_pairs,
        )

        # load node features
        if use_node_features:
            self.node_feature_tensor = torch.as_tensor(
                data=get_remapped_node_features(
                    dataset=dataset_name,
                    label_to_id=self.entity_to_id,
                    drive_label_to_id=drive_label_to_id,
                    emb=emb,
                    # TODO: Allow to provide this via constructor parameter?
                    cache_root=None,
                    force=False,
                ),
                dtype=torch.get_default_dtype(),
            )
            dim = self.node_feature_tensor.shape[-1]
            if feature_dim is not None and dim != feature_dim:
                logger.warning(
                    f"Using pre-trained features with dim {dim}, although feature_dim={feature_dim} was passed. "
                    f"To disable this warning pass feature_dim=None."
                )
        else:
            if feature_dim is None:
                raise ValueError("Must provide feature_dim if random features shall be used.")
            self.node_feature_tensor = torch.rand(
                self.num_entities,
                feature_dim,
            )
        self.feature_dim = self.node_feature_tensor.shape[-1]
        self.max_num_qualifier_pairs = max_num_qualifier_pairs

        # store non qualifier only entities
        self.non_qualifier_only_entities = self.mapped_statements[:, [0, 2]].unique()

    def _process_inverse_relations(self):
        relations = list(sorted(set([e for s in self.statements for e in s[1::2]])))

        # Check if the triples are inverted already
        relations_already_inverted = self._check_already_inverted_relations(relations)

        if self.create_inverse_triples or relations_already_inverted:
            self.create_inverse_triples = True
            if relations_already_inverted:
                logger.info(
                    f'Some triples already have suffix {INVERSE_SUFFIX}. '
                    f'Creating TriplesFactory based on inverse triples',
                )
                self.relation_to_inverse = {
                    re.sub('_inverse$', '', relation): f"{re.sub('_inverse$', '', relation)}{INVERSE_SUFFIX}"
                    for relation in relations
                }

            else:
                self.relation_to_inverse = {
                    relation: f"{relation}{INVERSE_SUFFIX}"
                    for relation in relations
                }
                inverse_statements = [[s[2], self.relation_to_inverse[s[1]], s[0], *s[3:]] for s in self.statements]
                # extend original triples with inverse ones
                self.statements.extend(inverse_statements)

        else:
            self.create_inverse_triples = False
            self.relation_to_inverse = None

    @staticmethod
    def _check_already_inverted_relations(relations: Iterable[str]) -> bool:
        for relation in relations:
            if relation.endswith(INVERSE_SUFFIX):
                # We can terminate the search after finding the first inverse occurrence
                return True

        return False

    def create_data_object(self) -> Data:
        """Create PyTorch Geometric data object."""
        node_features = self._create_node_feature_tensor()
        edge_index = self._create_edge_index()
        qualifier_index = self._create_qualifier_index()
        edge_type = self._create_edge_type()

        entities = torch.as_tensor(
            data=sorted(set(self.entity_to_id.values())),  # exclude padding entity
            dtype=torch.long,
        )

        return Data(
            x=node_features,
            edge_index=edge_index,
            qualifier_index=qualifier_index,
            edge_type=edge_type,
            entities=entities,
        )

    def _compose_qualifier_batch(self, keys, mapping) -> torch.LongTensor:
        # batch.shape:
        batch_size = len(keys)

        # Lookup IDs of qualifiers
        qualifier_ids = [
            mapping.get(tuple(k.tolist()))
            for k in keys
        ]

        # Determine maximum length (for dynamic padding)
        max_len = max(map(len, qualifier_ids))

        # bound maximum length for guaranteed max memory consumption
        # TODO: num_qualifier_pairs = max_num_qualifier_pairs ?
        max_len = min(max_len, self.max_num_qualifier_pairs)

        # Allocate result
        result = torch.full(size=(batch_size, max_len), fill_value=-1, dtype=torch.long)

        # Retrieve qualifiers
        for i, this_qualifier_ids in enumerate(qualifier_ids):
            # limit number of qualifiers
            # TODO: shuffle?
            this_qualifier_ids = this_qualifier_ids[:max_len]
            result[i, :len(this_qualifier_ids)] = torch.as_tensor(data=this_qualifier_ids, dtype=torch.long)

        return result

    def _create_qualifier_index(self):
        """
        Create a COO matrix of shape (3, num_qualifiers). Only non-zero (non-padded) qualifiers are retained.
        row0: qualifying relations
        row1: qualifying entities
        row2: index row which connects a pair (qual_r, qual_e) to a statement index k
        :return: shape: (3, num_qualifiers)
        """
        if self.max_num_qualifier_pairs is None or self.max_num_qualifier_pairs == 0:
            return None

        qual_relations, qual_entities, qual_k = [], [], []

        # It is assumed that statements are already padded
        for triple_id, statement in enumerate(self.mapped_statements):
            qualifiers = statement[3:]
            entities = qualifiers[1::2]
            relations = qualifiers[::2]
            # Ensure that PADDING has id=0
            non_zero_rels = relations[np.nonzero(relations)]
            non_zero_ents = entities[np.nonzero(entities)]
            assert len(non_zero_rels) == len(non_zero_ents), \
                "Number of non-padded qualifying relations is not equal to the # of qualifying entities"
            num_qualifier_pairs = non_zero_ents.shape[0]

            for j in range(num_qualifier_pairs):
                qual_relations.append(non_zero_rels[j].item())
                qual_entities.append(non_zero_ents[j].item())
                qual_k.append(triple_id)

        qualifier_index = torch.stack([
            torch.tensor(qual_relations, dtype=torch.long),
            torch.tensor(qual_entities, dtype=torch.long),
            torch.tensor(qual_k, dtype=torch.long)
        ], dim=0)

        if self.create_inverse_triples:
            # qualifier index is the same for inverse statements
            qualifier_index[2, len(qual_relations) // 2:] = qualifier_index[2, :len(qual_relations) // 2]

        return qualifier_index

    def _create_node_feature_tensor(self) -> torch.Tensor:
        """Create the node feature tensor."""
        if self.node_feature_tensor is not None:
            return self.node_feature_tensor
        if self.one_hot_encoding:
            return torch.eye(n=len(self.entity_to_id))
        else:
            return torch.empty(len(self.entity_to_id), self.feature_dim)

    def _create_edge_index(self) -> torch.Tensor:
        """Create edge index where first row represents the source nodes and the second row the target nodes."""
        mapped_heads = self.mapped_statements[:, 0].view(1, -1)
        mapped_tails = self.mapped_statements[:, 2].view(1, -1)
        edge_index = torch.cat([mapped_heads, mapped_tails], dim=0)
        return edge_index

    def _create_edge_type(self) -> torch.Tensor:
        """Create edge type tensor where each entry correspond to the relationship type of a triple in the dataset."""
        # Inverse triples are created in the base class

        return self.mapped_statements[:, 1]

    def _create_mappings_rdf_star(
        self,
        create_entity_to_id: bool = True,
        create_relation_to_id: bool = True,
    ):
        """."""

        max_length = self.max_num_qualifier_pairs * 2 + 3
        # Ensure that no entity/relation is part of the mappings that will not be seen
        statements = [s[:max_length] if len(s) > max_length else s for s in self.statements]

        if create_relation_to_id:
            relations = [PADDING] + list(sorted(set([r for s in statements for r in s[1::2]])))
            if self.create_inverse_triples:
                relations = [PADDING] + [elem for pair in
                                         zip(list(self.relation_to_inverse.keys()),
                                             list(self.relation_to_inverse.values()))
                                         for elem in pair]
                # old
                # relations = [PADDING] + list(self.relation_to_inverse.keys()) + (list(self.relation_to_inverse.values()))
            # When creating inverse relations, we will also generate padding_inverse which will never be used
            # unless we use score_r()
            self.relation_to_id = {relation: id for id, relation in enumerate(relations)}

        if create_entity_to_id:
            entities = [PADDING] + list(sorted(set([e for s in statements for e in s[::2]])))
            self.entity_to_id = {entity: id for id, entity in enumerate(entities)}

    def create_lcwa_instances(self, use_tqdm: Optional[bool] = None) -> LCWAInstances:
        """Create LCWA instances for this factory's statements."""
        s_p_q_to_multi_tails = _create_multi_label_tails_instance(
            mapped_statements=self.mapped_statements,
            use_tqdm=use_tqdm,
        )
        spq, multi_o = zip(*s_p_q_to_multi_tails.items())
        mapped_statements: torch.LongTensor = torch.tensor(spq, dtype=torch.long)
        labels = np.array([np.array(item) for item in multi_o], dtype=object)

        # create mask
        entity_mask = torch.zeros(self.num_entities, dtype=torch.bool)
        entity_mask[self.non_qualifier_only_entities] = True

        return LCWAInstances(
            mapped_statements=mapped_statements,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            labels=labels,
            # entity_mask=entity_mask,
        )

    def create_slcwa_instances(self) -> SLCWAInstances:
        """Create sLCWA instances for this factory's statements."""
        return SLCWAInstances(
            mapped_statements=self.mapped_statements,
            entity_to_id=self.entity_to_id,
            relation_to_id=self.relation_to_id,
            entities=self.non_qualifier_only_entities,
        )

    @property
    def num_entities(self) -> int:
        """The number of unique entities."""
        return len(self.entity_to_id)

    @property
    def num_relations(self) -> int:
        """The number of unique relations."""
        return len(self.relation_to_id)

    @property
    def statement_length(self) -> int:
        """The number of unique relations."""
        return self.max_num_qualifier_pairs * 2 + 3

    @property
    def num_statements(self) -> int:
        """The number of statements."""
        return self.mapped_statements.shape[0]

    @property
    def qualifier_ratio(self) -> float:
        """Return the percentage of statements with qualifiers."""
        return (~(self.mapped_statements[:, 3::2] == self.padding_idx).all(dim=1)).sum().item() / self.num_statements

    def extra_repr(self) -> str:
        """Extra representation."""
        return f"num_entities={self.num_entities:,}, " \
               f"num_relations={self.num_relations:,}, " \
               f"num_statements={self.num_statements:,}, " \
               f"max_num_qualifier_pairs={self.max_num_qualifier_pairs}, " \
               f"qualifier_ratio={self.qualifier_ratio:.2%}"

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self.extra_repr()})"

    # compatability with pykeen
    @property
    def num_triples(self) -> int:
        return self.num_statements


def _create_multi_label_tails_instance(
    mapped_statements: torch.Tensor,
    use_tqdm: Optional[bool] = None,
) -> Dict[Tuple[int, int], List[int]]:
    """Create for each (h,r,q*) pair the multi tail label."""
    logger.debug('Creating multi label tails instance')

    '''
    The mapped triples matrix has to be a numpy array to ensure correct pair hashing, as explained in
    https://github.com/pykeen/pykeen/commit/1bc71fe4eb2f24190425b0a4d0b9d6c7b9c4653a
    '''
    mapped_statements = mapped_statements.cpu().detach().numpy()

    s_p_q_to_multi_tails_new = _create_multi_label_instances(
        mapped_statements,
        element_1_index=0,
        element_2_index=1,
        label_index=2,
        use_tqdm=use_tqdm,
    )

    logger.debug('Created multi label tails instance')

    return s_p_q_to_multi_tails_new


def _create_multi_label_instances(
    mapped_statements: torch.Tensor,
    element_1_index: int,
    element_2_index: int,
    label_index: int,
    use_tqdm: Optional[bool] = None,
) -> Dict[Tuple[int, ...], List[int]]:
    """Create for each (element_1, element_2) pair the multi-label."""
    instance_to_multi_label = defaultdict(set)

    if use_tqdm is None:
        use_tqdm = True

    it = mapped_statements

    if use_tqdm:
        it = tqdm(mapped_statements, unit='statement', unit_scale=True, desc='Grouping statements')

    for row in it:
        instance = tuple([row[element_1_index], row[element_2_index]] + row[3:].tolist())
        instance_to_multi_label[instance].add(
            row[label_index])

    # Create lists out of sets for proper numpy indexing when loading the labels
    # TODO is there a need to have a canonical sort order here?
    instance_to_multi_label_new = {
        key: list(value)
        for key, value in instance_to_multi_label.items()
    }

    return instance_to_multi_label_new


def get_semi_inductive_factories(
    dataset_name: str,
    max_num_qualifier_pairs: Optional[int] = None,
    create_inverse_triples: bool = True,
) -> Tuple[StatementsFactory, StatementsFactory, StatementsFactory]:
    """Get factories for semi-inductive split."""
    dataset = dataset_resolver.make(query=dataset_name)
    paths = {
        key: dataset.get_path(split="semi_inductive", part=key)
        for key in [
            "train",
            "validation",
            "test",
        ]
    }

    # Step 1: Get relation_to_id based on relations in train
    tmp_stmt_factory = StatementsFactory(
        path=paths['train'],
        dataset_name=dataset_name,
        use_node_features=True,
        is_inductive=False,
        enforce_node_features=True,
        create_inverse_triples=create_inverse_triples,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )

    relation_to_id = tmp_stmt_factory.relation_to_id

    # Step 2: Create entity_to_id based on all entities. Entities need to be connected to train-relations
    tmp_stmt_factory = StatementsFactory(
        path=list(paths.values()),
        dataset_name=dataset_name,
        create_inverse_triples=create_inverse_triples,
        relation_to_id=relation_to_id,
        use_node_features=True,
        is_inductive=True,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Full: {tmp_stmt_factory}")

    entity_to_id = tmp_stmt_factory.entity_to_id

    del tmp_stmt_factory

    # Create statement factories
    train_sf = StatementsFactory(
        path=paths["train"],
        dataset_name=dataset_name,
        entity_to_id=entity_to_id,
        relation_to_id=relation_to_id,
        create_inverse_triples=create_inverse_triples,
        use_node_features=True,
        is_inductive=False,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Training: {train_sf}")

    validation_sf = StatementsFactory(
        path=paths["validation"],
        dataset_name=dataset_name,
        is_inductive=True,
        create_inverse_triples=False,
        # inductive entities
        entity_to_id=train_sf.entity_to_id,
        # shared relations
        relation_to_id=train_sf.relation_to_id,
        use_node_features=True,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Validation: {validation_sf}")

    # In fully inductive setting, validation graph consists of training entities
    test_sf = StatementsFactory(
        path=paths["test"],
        dataset_name=dataset_name,
        create_inverse_triples=False,
        is_inductive=True,
        # inductive entities
        entity_to_id=train_sf.entity_to_id,
        # inductive relations
        relation_to_id=train_sf.relation_to_id,
        use_node_features=True,
        enforce_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    logging.info(f"Testing: {test_sf}")

    return train_sf, validation_sf, test_sf


def get_fully_inductive_factories(
    dataset_name: str,
    dataset_version: str,
    max_num_qualifier_pairs: int,
    create_inverse_triples: bool,
) -> Tuple[StatementsFactory, StatementsFactory, StatementsFactory, StatementsFactory]:
    """Create statement factories for the fully inductive setting."""
    dataset = dataset_resolver.make(query=dataset_name)
    paths = {
        key: dataset.get_path(split="inductive", part=key, version=dataset_version)
        for key in [
            "train",
            "inference",
            "validation",
            "test",
        ]
    }
    train_transductive_sf = StatementsFactory(
        path=paths["train"],
        dataset_name=dataset_name,
        create_inverse_triples=create_inverse_triples,
        use_node_features=True,
        is_inductive=False,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    train_inductive_sf = StatementsFactory(
        path=paths["inference"],
        dataset_name=dataset_name,
        create_inverse_triples=create_inverse_triples,
        relation_to_id=train_transductive_sf.relation_to_id,
        is_inductive=True,
        use_node_features=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    # In fully inductive setting, validation graph consists of training entities
    validation_sf = StatementsFactory(
        path=paths["validation"],
        dataset_name=dataset_name,
        entity_to_id=train_inductive_sf.entity_to_id,
        relation_to_id=train_inductive_sf.relation_to_id,
        use_node_features=True,
        is_inductive=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    # Test graph contains only unseen entities, therefore, create
    test_sf = StatementsFactory(
        path=paths["test"],
        dataset_name=dataset_name,
        entity_to_id=train_inductive_sf.entity_to_id,
        relation_to_id=train_inductive_sf.relation_to_id,
        use_node_features=True,
        is_inductive=True,
        max_num_qualifier_pairs=max_num_qualifier_pairs,
    )
    return train_transductive_sf, train_inductive_sf, test_sf, validation_sf


# buffer loading (for HPO)
@functools.lru_cache(maxsize=1)
def get_factories(
    dataset_name: str,
    dataset_version: Optional[str],
    inductive_setting: str,
    max_num_qualifier_pairs: int,
    create_inverse_triples: bool,
) -> Tuple[StatementsFactory, StatementsFactory, StatementsFactory, StatementsFactory]:
    """Load triples factories."""
    if inductive_setting == "semi":
        train_sf, validation_sf, test_sf = get_semi_inductive_factories(
            dataset_name=dataset_name,
            max_num_qualifier_pairs=max_num_qualifier_pairs,
            create_inverse_triples=create_inverse_triples,
        )
        inference_sf = train_sf

    elif inductive_setting == "full":
        train_sf, inference_sf, test_sf, validation_sf = get_fully_inductive_factories(
            dataset_name=dataset_name,
            dataset_version=dataset_version,
            max_num_qualifier_pairs=max_num_qualifier_pairs,
            create_inverse_triples=create_inverse_triples,
        )

    else:
        raise ValueError(f"Unknown inductive setting: {inductive_setting}")

    return inference_sf, test_sf, train_sf, validation_sf

# -*- coding: utf-8 -*- A graph factory that creates the PyTorchGeometric data object, and holds references to the textual descriptions. # Map to id # Pad statement Create graph and save it as torch_geometric.data. Initialize the factory.          :param path:             path to the graph file         :param dataset_name:             The dataset name.         :param statements:             directly input triples (optional)         :param entity_to_id:             mapping from labels to a numeric id. If None, infer from statements.         :param relation_to_id:             mapping relation type to an numeric id. If None, infer from statements.         :param use_node_features:             whether to use node features or generate random ones.         :param feature_dim:             dimension of random features for nodes and edges. # Filter statements based on node features # convert to tuple # if head or tail have no features -> skip statement # if any relation is unknown: # filter qualifiers by entities # filter qualifiers by relations # Generate entity mapping if necessary # Note: Mapped statements are not padded # load node features # TODO: Allow to provide this via constructor parameter? # store non qualifier only entities # Check if the triples are inverted already # extend original triples with inverse ones # We can terminate the search after finding the first inverse occurrence Create PyTorch Geometric data object. # exclude padding entity # batch.shape: # Lookup IDs of qualifiers # Determine maximum length (for dynamic padding) # bound maximum length for guaranteed max memory consumption # TODO: num_qualifier_pairs = max_num_qualifier_pairs ? # Allocate result # Retrieve qualifiers # limit number of qualifiers # TODO: shuffle? Create a COO matrix of shape (3, num_qualifiers). Only non-zero (non-padded) qualifiers are retained.         row0: qualifying relations         row1: qualifying entities         row2: index row which connects a pair (qual_r, qual_e) to a statement index k         :return: shape: (3, num_qualifiers) # It is assumed that statements are already padded # Ensure that PADDING has id=0 # of qualifying entities" # qualifier index is the same for inverse statements Create the node feature tensor. Create edge index where first row represents the source nodes and the second row the target nodes. Create edge type tensor where each entry correspond to the relationship type of a triple in the dataset. # Inverse triples are created in the base class . # Ensure that no entity/relation is part of the mappings that will not be seen # old # relations = [PADDING] + list(self.relation_to_inverse.keys()) + (list(self.relation_to_inverse.values())) # When creating inverse relations, we will also generate padding_inverse which will never be used # unless we use score_r() Create LCWA instances for this factory's statements. # create mask # entity_mask=entity_mask, Create sLCWA instances for this factory's statements. The number of unique entities. The number of unique relations. The number of unique relations. The number of statements. Return the percentage of statements with qualifiers. Extra representation. # compatability with pykeen Create for each (h,r,q*) pair the multi tail label. The mapped triples matrix has to be a numpy array to ensure correct pair hashing, as explained in     https://github.com/pykeen/pykeen/commit/1bc71fe4eb2f24190425b0a4d0b9d6c7b9c4653a Create for each (element_1, element_2) pair the multi-label. # Create lists out of sets for proper numpy indexing when loading the labels # TODO is there a need to have a canonical sort order here? Get factories for semi-inductive split. # Step 1: Get relation_to_id based on relations in train # Step 2: Create entity_to_id based on all entities. Entities need to be connected to train-relations # Create statement factories # inductive entities # shared relations # In fully inductive setting, validation graph consists of training entities # inductive entities # inductive relations Create statement factories for the fully inductive setting. # In fully inductive setting, validation graph consists of training entities # Test graph contains only unseen entities, therefore, create # buffer loading (for HPO) Load triples factories.