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

"""Implementation the command line interface needed for Structured Model (SM)."""

from collections import OrderedDict

import pykeen.constants as pkc
from pykeen.cli.utils.cli_print_msg_helper import (
    print_batch_size_message, print_embedding_dimension_info_message, print_entity_normalization_message,
    print_hpo_batch_sizes_message, print_hpo_embedding_dimensions_message, print_hpo_entity_normalization_norms_message,
    print_hpo_epochs_message, print_hpo_learning_rates_message, print_hpo_margin_losses_message,
    print_hpo_scoring_fcts_message, print_learning_rate_message, print_number_epochs_message, print_scoring_fct_message,
    print_section_divider, print_training_embedding_dimension_message, print_training_margin_loss_message,
)
from pykeen.cli.utils.cli_query_helper import (
    select_float_value, select_float_values, select_integer_value, select_positive_integer_values,
)
from pykeen.utilities.summarize import get_config_dict

__all__ = [
    'configure_se_hpo_pipeline',
    'configure_se_training_pipeline',
]


def configure_se_training_pipeline(model_name: str):
    config = OrderedDict()
    config[pkc.KG_EMBEDDING_MODEL_NAME] = model_name

    # Step 1: Query embedding dimension
    print_training_embedding_dimension_message()
    print_embedding_dimension_info_message()
    embedding_dimension = select_integer_value(
        print_msg=pkc.EMBEDDING_DIMENSION_PRINT_MSG,
        prompt_msg=pkc.EMBEDDING_DIMENSION_PROMPT_MSG,
        error_msg=pkc.EMBEDDING_DIMENSION_ERROR_MSG,
    )
    config[pkc.EMBEDDING_DIM] = embedding_dimension
    print_section_divider()

    # Step 2: Query margin loss
    print_training_margin_loss_message()
    magin_loss = select_float_value(
        print_msg=pkc.MARGIN_LOSS_PRINT_MSG,
        prompt_msg=pkc.MARGIN_LOSS_PROMPT_MSG,
        error_msg=pkc.MARGIN_LOSS_ERROR_MSG,
    )
    config[pkc.MARGIN_LOSS] = magin_loss
    print_section_divider()

    # Step 3: Query L_p norm as scoring function
    print_scoring_fct_message()
    scoring_fct_norm = select_integer_value(
        print_msg=pkc.NORM_SCORING_FUNCTION_PRINT_MSG,
        prompt_msg=pkc.NORM_SCORING_FUNCTION_PROMPT_MSG,
        error_msg=pkc.NORM_SCORING_FUNCTION_ERROR_MSG,
    )
    config[pkc.SCORING_FUNCTION_NORM] = scoring_fct_norm
    print_section_divider()

    # Step 4: Query L_p norm for normalizing the entities
    print_entity_normalization_message()
    entity_normalization_norm = select_integer_value(
        print_msg=pkc.ENTITIES_NORMALIZATION_PRINT_MSG,
        prompt_msg=pkc.ENTITIES_NORMALIZATION_PROMPT_MSG,
        error_msg=pkc.ENTITIES_NORMALIZATION_ERROR_MSG,
    )
    config[pkc.NORM_FOR_NORMALIZATION_OF_ENTITIES] = entity_normalization_norm
    print_section_divider()

    # Step 5: Query learning rate
    print_learning_rate_message()
    learning_rate = select_float_value(
        print_msg=pkc.LEARNING_RATE_PRINT_MSG,
        prompt_msg=pkc.LEARNING_RATE_PROMPT_MSG,
        error_msg=pkc.LEARNING_RATE_ERROR_MSG,
    )
    config[pkc.LEARNING_RATE] = learning_rate
    print_section_divider()

    # Step 6: Query batch size
    print_batch_size_message()
    batch_size = select_integer_value(
        print_msg=pkc.BATCH_SIZE_PRINT_MSG,
        prompt_msg=pkc.BATCH_SIZE_PROMPT_MSG,
        error_msg=pkc.BATCH_SIZE_ERROR_MSG,
    )
    config[pkc.BATCH_SIZE] = batch_size
    print_section_divider()

    # Step 7: Query number of epochs
    print_number_epochs_message()
    number_epochs = select_integer_value(
        print_msg=pkc.EPOCH_PRINT_MSG,
        prompt_msg=pkc.EPOCH_PROMPT_MSG,
        error_msg=pkc.EPOCH_ERROR_MSG,
    )
    config[pkc.NUM_EPOCHS] = number_epochs
    print_section_divider()

    return config


def configure_se_hpo_pipeline(model_name):
    config = get_config_dict(model_name)

    # Step 1: Query embedding dimensions
    print_hpo_embedding_dimensions_message()
    embedding_dimensions = select_positive_integer_values(
        pkc.EMBEDDING_DIMENSIONS_PRINT_MSG,
        pkc.EMBEDDING_DIMENSIONS_PROMPT_MSG,
        pkc.EMBEDDING_DIMENSIONS_ERROR_MSG,
    )
    config[pkc.EMBEDDING_DIM] = embedding_dimensions
    print_section_divider()

    # Step 2: Query margin loss
    print_hpo_margin_losses_message()
    magin_loss = select_float_values(
        print_msg=pkc.MARGIN_LOSSES_PRINT_MSG,
        prompt_msg=pkc.MARGIN_LOSSES_PROMPT_MSG,
        error_msg=pkc.MARGIN_LOSSES_ERROR_MSG,
    )

    config[pkc.MARGIN_LOSS] = magin_loss
    print_section_divider()

    # Step 3: Query L_p norms to use as scoring function
    print_hpo_scoring_fcts_message()
    scoring_fct_norm = select_positive_integer_values(
        print_msg=pkc.NORMS_SCORING_FUNCTION_PRINT_MSG,
        prompt_msg=pkc.NORMS_SCORING_FUNCTION_PROMPT_MSG,
        error_msg=pkc.NORMS_SCORING_FUNCTION_ERROR_MSG,
    )
    config[pkc.SCORING_FUNCTION_NORM] = scoring_fct_norm
    print_section_divider()

    # Step 4: Query L_p norms for normalizing the entities
    print_hpo_entity_normalization_norms_message()
    entity_normalization_norm = select_positive_integer_values(
        print_msg=pkc.NORMS_FOR_NORMALIZATION_OF_ENTITIES_PRINT_MSG,
        prompt_msg=pkc.NORMS_FOR_NORMALIZATION_OF_ENTITIES_PROMPT_MSG,
        error_msg=pkc.NORMS_FOR_NORMALIZATION_OF_ENTITIES_ERROR_MSG,
    )
    config[pkc.NORM_FOR_NORMALIZATION_OF_ENTITIES] = entity_normalization_norm
    print_section_divider()

    # Step 5: Query learning rate
    print_hpo_learning_rates_message()
    learning_rate = select_float_values(
        print_msg=pkc.LEARNING_RATES_PRINT_MSG,
        prompt_msg=pkc.LEARNING_RATES_PROMPT_MSG,
        error_msg=pkc.LEARNING_RATES_ERROR_MSG,
    )
    config[pkc.LEARNING_RATE] = learning_rate
    print_section_divider()

    # Step 6: Query batch size
    print_hpo_batch_sizes_message()
    batch_size = select_positive_integer_values(
        print_msg=pkc.BATCH_SIZES_PRINT_MSG,
        prompt_msg=pkc.BATCH_SIZES_PROMPT_MSG,
        error_msg=pkc.BATCH_SIZES_ERROR_MSG,
    )
    config[pkc.BATCH_SIZE] = batch_size
    print_section_divider()

    # Step 7: Query number of epochs
    print_hpo_epochs_message()
    number_epochs = select_positive_integer_values(
        print_msg=pkc.EPOCHS_PRINT_MSG,
        prompt_msg=pkc.EPOCHS_PROMPT_MSG,
        error_msg=pkc.EPOCHS_ERROR_MSG,
    )
    config[pkc.NUM_EPOCHS] = number_epochs
    print_section_divider()

    return config
