<reponame>Nudin/PyKEEN
# -*- coding: utf-8 -*-

"""A hyper-parameter optimizer that uses random search."""

import random
from collections import OrderedDict
from typing import Any, Dict, List, Optional

import numpy as np
import pykeen.constants as pkc
import torch
from pykeen.hpo.utils import HPOptimizer, HPOptimizerResult
from pykeen.kge_models import get_kge_model
from pykeen.utilities.evaluation_utils.metrics_computations import (
    compute_metric_results,
)
from torch.nn import Module
from tqdm import trange

__all__ = ["RandomSearch"]


class RandomSearch(HPOptimizer):
    """A hyper-parameter optimizer that uses random search."""

    def _sample_conv_e_params(self, hyperparams_dict) -> Dict[str, Any]:
        kg_model_config = OrderedDict()
        # Sample params which are dependent on each other

        hyperparams_dict = hyperparams_dict.copy()
        embedding_dimensions = hyperparams_dict[pkc.EMBEDDING_DIM]
        sampled_index = random.choice(range(len(embedding_dimensions)))
        kg_model_config[pkc.EMBEDDING_DIM] = hyperparams_dict[pkc.EMBEDDING_DIM][
            sampled_index
        ]
        kg_model_config[pkc.CONV_E_HEIGHT] = hyperparams_dict[pkc.CONV_E_HEIGHT][
            sampled_index
        ]
        kg_model_config[pkc.CONV_E_WIDTH] = hyperparams_dict[pkc.CONV_E_WIDTH][
            sampled_index
        ]
        kg_model_config[pkc.CONV_E_KERNEL_HEIGHT] = hyperparams_dict[
            pkc.CONV_E_KERNEL_HEIGHT
        ][sampled_index]
        kg_model_config[pkc.CONV_E_KERNEL_WIDTH] = hyperparams_dict[
            pkc.CONV_E_KERNEL_WIDTH
        ][sampled_index]

        del hyperparams_dict[pkc.EMBEDDING_DIM]
        del hyperparams_dict[pkc.CONV_E_HEIGHT]
        del hyperparams_dict[pkc.CONV_E_WIDTH]
        del hyperparams_dict[pkc.CONV_E_KERNEL_HEIGHT]
        del hyperparams_dict[pkc.CONV_E_KERNEL_WIDTH]

        kg_model_config.update(self._sample_parameter_value(hyperparams_dict))

        return kg_model_config

    def optimize_hyperparams(
        self,
        mapped_train_triples,
        mapped_test_triples,
        entity_to_id,
        rel_to_id,
        config,
        device,
        seed: Optional[int] = None,
        k_evaluation: int = 10,
    ) -> HPOptimizerResult:
        if seed is not None:
            torch.manual_seed(config[pkc.SEED])

        trained_kge_models: List[Module] = []
        epoch_losses: List[List[float]] = []
        hits_at_k_evaluations: List[float] = []
        entity_to_ids: List[Dict[int, str]] = []
        rel_to_ids: List[Dict[int, str]] = []
        models_params: List[Dict] = []
        eval_summaries: List = []

        config = config.copy()
        max_iters = config[pkc.NUM_OF_HPO_ITERS]

        sample_fct = (
            self._sample_conv_e_params
            if config[pkc.KG_EMBEDDING_MODEL_NAME] == pkc.CONV_E_NAME
            else self._sample_parameter_value
        )

        for _ in trange(max_iters, desc="HPO Iteration"):
            # Sample hyper-params
            kge_model_config: Dict[str, Any] = sample_fct(config)
            kge_model_config[pkc.SEED]: int = seed
            kge_model_config[
                pkc.PREFERRED_DEVICE
            ]: str = pkc.CPU if device == pkc.CPU else pkc.GPU

            # Configure defined model
            kge_model: Module = get_kge_model(config=kge_model_config)

            # Load class params
            kge_model.entity_label_to_id: Dict[str, int] = entity_to_id
            kge_model.relation_label_to_id: Dict[str, int] = rel_to_id
            kge_model.num_entities: int = len(entity_to_id)
            kge_model.num_relations: int = len(rel_to_id)

            models_params.append(kge_model_config)
            entity_to_ids.append(entity_to_id)
            rel_to_ids.append(rel_to_id)

            batch_size = kge_model_config[pkc.BATCH_SIZE]
            num_epochs = kge_model_config[pkc.NUM_EPOCHS]
            learning_rate = kge_model_config[pkc.LEARNING_RATE]

            epoch_loss = kge_model.fit(
                pos_triples=mapped_train_triples,
                learning_rate=learning_rate,
                num_epochs=num_epochs,
                batch_size=batch_size,
                tqdm_kwargs=dict(leave=False),
            )

            trained_kge_model = kge_model

            # Evaluate trained model
            metric_results = compute_metric_results(
                kg_embedding_model=trained_kge_model,
                mapped_train_triples=mapped_train_triples,
                mapped_test_triples=mapped_test_triples,
                device=device,
            )

            # TODO: Define HPO metric
            eval_summaries.append(metric_results)

            trained_kge_models.append(trained_kge_model)
            epoch_losses.append(epoch_loss)

            hits_at_k_evaluation = metric_results.hits_at_k[k_evaluation]
            hits_at_k_evaluations.append(hits_at_k_evaluation)

        index_of_max = int(np.argmax(a=hits_at_k_evaluations))

        return (
            trained_kge_models[index_of_max],
            epoch_losses[index_of_max],
            entity_to_ids[index_of_max],
            rel_to_ids[index_of_max],
            eval_summaries[index_of_max],
            models_params[index_of_max],
        )

    @classmethod
    def run(
        cls,
        mapped_train_triples: np.ndarray,
        mapped_test_triples: np.ndarray,
        entity_to_id: Dict[int, str],
        rel_to_id: Dict[int, str],
        config: Dict,
        device,
        seed,
    ) -> HPOptimizerResult:
        return cls().optimize_hyperparams(
            mapped_train_triples=mapped_train_triples,
            mapped_test_triples=mapped_test_triples,
            entity_to_id=entity_to_id,
            rel_to_id=rel_to_id,
            config=config,
            device=device,
            seed=seed,
        )

# -*- coding: utf-8 -*-

"""A hyper-parameter optimizer that uses random search."""

import random
from collections import OrderedDict
from typing import Any, Dict, List, Optional

import numpy as np
import pykeen.constants as pkc
import torch
from pykeen.hpo.utils import HPOptimizer, HPOptimizerResult
from pykeen.kge_models import get_kge_model
from pykeen.utilities.evaluation_utils.metrics_computations import (
    compute_metric_results,
)
from torch.nn import Module
from tqdm import trange

__all__ = ["RandomSearch"]


class RandomSearch(HPOptimizer):
    """A hyper-parameter optimizer that uses random search."""

    def _sample_conv_e_params(self, hyperparams_dict) -> Dict[str, Any]:
        kg_model_config = OrderedDict()
        # Sample params which are dependent on each other

        hyperparams_dict = hyperparams_dict.copy()
        embedding_dimensions = hyperparams_dict[pkc.EMBEDDING_DIM]
        sampled_index = random.choice(range(len(embedding_dimensions)))
        kg_model_config[pkc.EMBEDDING_DIM] = hyperparams_dict[pkc.EMBEDDING_DIM][
            sampled_index
        ]
        kg_model_config[pkc.CONV_E_HEIGHT] = hyperparams_dict[pkc.CONV_E_HEIGHT][
            sampled_index
        ]
        kg_model_config[pkc.CONV_E_WIDTH] = hyperparams_dict[pkc.CONV_E_WIDTH][
            sampled_index
        ]
        kg_model_config[pkc.CONV_E_KERNEL_HEIGHT] = hyperparams_dict[
            pkc.CONV_E_KERNEL_HEIGHT
        ][sampled_index]
        kg_model_config[pkc.CONV_E_KERNEL_WIDTH] = hyperparams_dict[
            pkc.CONV_E_KERNEL_WIDTH
        ][sampled_index]

        del hyperparams_dict[pkc.EMBEDDING_DIM]
        del hyperparams_dict[pkc.CONV_E_HEIGHT]
        del hyperparams_dict[pkc.CONV_E_WIDTH]
        del hyperparams_dict[pkc.CONV_E_KERNEL_HEIGHT]
        del hyperparams_dict[pkc.CONV_E_KERNEL_WIDTH]

        kg_model_config.update(self._sample_parameter_value(hyperparams_dict))

        return kg_model_config

    def optimize_hyperparams(
        self,
        mapped_train_triples,
        mapped_test_triples,
        entity_to_id,
        rel_to_id,
        config,
        device,
        seed: Optional[int] = None,
        k_evaluation: int = 10,
    ) -> HPOptimizerResult:
        if seed is not None:
            torch.manual_seed(config[pkc.SEED])

        trained_kge_models: List[Module] = []
        epoch_losses: List[List[float]] = []
        hits_at_k_evaluations: List[float] = []
        entity_to_ids: List[Dict[int, str]] = []
        rel_to_ids: List[Dict[int, str]] = []
        models_params: List[Dict] = []
        eval_summaries: List = []

        config = config.copy()
        max_iters = config[pkc.NUM_OF_HPO_ITERS]

        sample_fct = (
            self._sample_conv_e_params
            if config[pkc.KG_EMBEDDING_MODEL_NAME] == pkc.CONV_E_NAME
            else self._sample_parameter_value
        )

        for _ in trange(max_iters, desc="HPO Iteration"):
            # Sample hyper-params
            kge_model_config: Dict[str, Any] = sample_fct(config)
            kge_model_config[pkc.SEED]: int = seed
            kge_model_config[
                pkc.PREFERRED_DEVICE
            ]: str = pkc.CPU if device == pkc.CPU else pkc.GPU

            # Configure defined model
            kge_model: Module = get_kge_model(config=kge_model_config)

            # Load class params
            kge_model.entity_label_to_id: Dict[str, int] = entity_to_id
            kge_model.relation_label_to_id: Dict[str, int] = rel_to_id
            kge_model.num_entities: int = len(entity_to_id)
            kge_model.num_relations: int = len(rel_to_id)

            models_params.append(kge_model_config)
            entity_to_ids.append(entity_to_id)
            rel_to_ids.append(rel_to_id)

            batch_size = kge_model_config[pkc.BATCH_SIZE]
            num_epochs = kge_model_config[pkc.NUM_EPOCHS]
            learning_rate = kge_model_config[pkc.LEARNING_RATE]

            epoch_loss = kge_model.fit(
                pos_triples=mapped_train_triples,
                learning_rate=learning_rate,
                num_epochs=num_epochs,
                batch_size=batch_size,
                tqdm_kwargs=dict(leave=False),
            )

            trained_kge_model = kge_model

            # Evaluate trained model
            metric_results = compute_metric_results(
                kg_embedding_model=trained_kge_model,
                mapped_train_triples=mapped_train_triples,
                mapped_test_triples=mapped_test_triples,
                device=device,
            )

            # TODO: Define HPO metric
            eval_summaries.append(metric_results)

            trained_kge_models.append(trained_kge_model)
            epoch_losses.append(epoch_loss)

            hits_at_k_evaluation = metric_results.hits_at_k[k_evaluation]
            hits_at_k_evaluations.append(hits_at_k_evaluation)

        index_of_max = int(np.argmax(a=hits_at_k_evaluations))

        return (
            trained_kge_models[index_of_max],
            epoch_losses[index_of_max],
            entity_to_ids[index_of_max],
            rel_to_ids[index_of_max],
            eval_summaries[index_of_max],
            models_params[index_of_max],
        )

    @classmethod
    def run(
        cls,
        mapped_train_triples: np.ndarray,
        mapped_test_triples: np.ndarray,
        entity_to_id: Dict[int, str],
        rel_to_id: Dict[int, str],
        config: Dict,
        device,
        seed,
    ) -> HPOptimizerResult:
        return cls().optimize_hyperparams(
            mapped_train_triples=mapped_train_triples,
            mapped_test_triples=mapped_test_triples,
            entity_to_id=entity_to_id,
            rel_to_id=rel_to_id,
            config=config,
            device=device,
            seed=seed,
        )
# -*- coding: utf-8 -*- A hyper-parameter optimizer that uses random search. A hyper-parameter optimizer that uses random search. # Sample params which are dependent on each other # Sample hyper-params # Configure defined model # Load class params # Evaluate trained model # TODO: Define HPO metric