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

"""PyKEEN's command line interface."""

import json
import os
from collections import OrderedDict
from typing import Dict

import click

from pykeen.constants import (
    EXECUTION_MODE, FILTER_NEG_TRIPLES, HPO_ITERS_ERROR_MSG, HPO_ITERS_PRINT_MSG, HPO_ITERS_PROMPT_MSG, HPO_MODE,
    NUM_OF_HPO_ITERS, OUTPUT_DIREC, PREFERRED_DEVICE, PYKEEN, SEED, SEED_ERROR_MSG, SEED_PRINT_MSG, SEED_PROMPT_MSG,
    TEST_FILE_ERROR_MSG, TEST_FILE_PROMPT_MSG, TEST_SET_PATH, TEST_SET_RATIO, TRAINING_FILE_ERROR_MSG,
    TRAINING_FILE_PROMPT_MSG, TRAINING_MODE, TRAINING_SET_PATH,
)
from pykeen.predict import start_predictions_pipeline
from pykeen.run import run
from pykeen.utilities.cli_utils.cli_print_msg_helper import (
    print_ask_for_evlauation_message, print_execution_mode_message, print_filter_negative_triples_message, print_intro,
    print_model_selection_message, print_random_seed_message, print_section_divider, print_test_ratio_message,
    print_test_set_message, print_training_set_message, print_welcome_message,
)
from pykeen.utilities.cli_utils.cli_query_helper import (
    ask_for_evaluation, ask_for_filtering_of_negatives, ask_for_test_set, get_input_path, query_output_directory,
    select_embedding_model, select_integer_value, select_keen_execution_mode, select_preferred_device,
    select_ratio_for_test_set,
)
from pykeen.utilities.cli_utils.dicts import MODEL_HPO_CONFIG_FUNCS, MODEL_TRAINING_CONFIG_FUNCS
from pykeen.utilities.cli_utils.utils import summarize_results


def _configure_training_pipeline(model_name: str):
    model_config_func = MODEL_TRAINING_CONFIG_FUNCS.get(model_name)
    if model_config_func is None:
        raise KeyError(f'invalid model given: {model_name}')
    config = model_config_func(model_name)
    if config is None:
        raise NotImplementedError(f'{model_name} has not yet been implemented')
    return config


def _configure_hpo_pipeline(model_name: str):
    model_config_func = MODEL_HPO_CONFIG_FUNCS.get(model_name)
    if model_config_func is None:
        raise KeyError(f'invalid model given: {model_name}')
    config = model_config_func(model_name)
    if config is None:
        raise NotImplementedError(f'{model_name} has not yet been implemented')
    return config


def prompt_evaluation_parameters(config: Dict) -> None:
    """Prompt the user for evaluation parameters absed on the execution mode."""
    if config[EXECUTION_MODE] == TRAINING_MODE:
        # Step 1: Ask whether to evaluate the model
        print_ask_for_evlauation_message()
        is_evaluation_mode = ask_for_evaluation()
        print_section_divider()
    else:
        is_evaluation_mode = True

    # Step 2: Specify test set, if is_evaluation_mode==True
    if is_evaluation_mode:
        print_test_set_message()
        provide_test_set = ask_for_test_set()
        print_section_divider()

        if provide_test_set:
            test_set_path = get_input_path(
                prompt_msg=TEST_FILE_PROMPT_MSG,
                error_msg=TEST_FILE_ERROR_MSG,
            )
            config[TEST_SET_PATH] = test_set_path
        else:
            print_test_ratio_message()
            test_set_ratio = select_ratio_for_test_set()
            config[TEST_SET_RATIO] = test_set_ratio

        print_section_divider()

        # Ask whether to use filtered negative triples
        print_filter_negative_triples_message()
        filter_negative_triples = ask_for_filtering_of_negatives()
        config[FILTER_NEG_TRIPLES] = filter_negative_triples
        print_section_divider()


def prompt_training_file(config: Dict) -> None:
    """Prompt the user for a training file."""
    print_training_set_message()
    config[TRAINING_SET_PATH] = get_input_path(
        prompt_msg=TRAINING_FILE_PROMPT_MSG,
        error_msg=TRAINING_FILE_ERROR_MSG,
        is_dataset=True,
    )


def prompt_execution_mode(config: Dict, lib_name: str = PYKEEN) -> None:
    """Prompt the user for the execution mode."""
    print_execution_mode_message()
    config[EXECUTION_MODE] = select_keen_execution_mode(lib_name=lib_name)


def prompt_embedding_model() -> str:
    """Prompt the user to select an embedding model."""
    print_model_selection_message()
    model_name = select_embedding_model()
    print_section_divider()
    return model_name


def prompt_execution_parameters(config: Dict, model_name: str) -> None:
    """Prompt the user for execution mode parameters."""
    pykeen_exec_mode = config[EXECUTION_MODE]

    if pykeen_exec_mode == TRAINING_MODE:
        config.update(_configure_training_pipeline(model_name))

    elif pykeen_exec_mode == HPO_MODE:
        config.update(_configure_hpo_pipeline(model_name))

        # Query number of HPO iterations
        hpo_iter = select_integer_value(
            print_msg=HPO_ITERS_PRINT_MSG,
            prompt_msg=HPO_ITERS_PROMPT_MSG,
            error_msg=HPO_ITERS_ERROR_MSG,
        )
        config[NUM_OF_HPO_ITERS] = hpo_iter
        print_section_divider()


def prompt_random_seed(config):
    """Query random seed."""
    print_random_seed_message()
    config[SEED] = select_integer_value(
        print_msg=SEED_PRINT_MSG,
        prompt_msg=SEED_PROMPT_MSG,
        error_msg=SEED_ERROR_MSG,
    )


def prompt_device(config: Dict) -> None:
    """Prompt the user for their preferred evaluation device."""
    config[PREFERRED_DEVICE] = select_preferred_device()


def prompt_output_directory(config: Dict) -> None:
    """Prompt the user for the output directory."""
    config[OUTPUT_DIREC] = query_output_directory()


def prompt_config() -> Dict:
    """Prompt the user for the run configuration."""
    cfg = OrderedDict()

    # Step 1: Welcome + Intro
    print_welcome_message()
    print_section_divider()

    print_intro()
    print_section_divider()

    # Step 2: Ask for training file
    prompt_training_file(cfg)
    print_section_divider()

    # Step 3: Ask for execution mode
    prompt_execution_mode(cfg)
    print_section_divider()

    # Step 4: Ask for model
    model_name = prompt_embedding_model()
    print_section_divider()

    # Step 5: Query parameters depending on the selected execution mode
    prompt_execution_parameters(cfg, model_name=model_name)
    print_section_divider()

    # Step 5.5: Prompt for evaluation parameters depending on the selected execution mode
    prompt_evaluation_parameters(cfg)

    # Step 6: Please select a random seed
    prompt_random_seed(cfg)
    print_section_divider()

    # Step 7: Query device to train on
    prompt_device(cfg)
    print_section_divider()

    # Step 8: Define output directory
    prompt_output_directory(cfg)
    print_section_divider()

    return cfg


@click.command()
@click.option('-c', '--config', type=click.File(), help='A PyKEEN JSON configuration file')
def main(config):
    """PyKEEN: A software for training and evaluating knowledge graph embeddings."""
    if config is not None:
        config = json.load(config)
    else:
        config = prompt_config()

    run(config)


@click.command()
@click.option('-m', '--model-directory', type=click.Path(file_okay=False, dir_okay=True))
@click.option('-d', '--data-directory', type=click.Path(file_okay=False, dir_okay=True))
def predict(model_directory: str, data_directory: str):
    """Predict new links based on trained model."""
    start_predictions_pipeline(model_directory, data_directory)


@click.command()
@click.option('-d', '--directory', type=click.Path(file_okay=False, dir_okay=True), default=os.getcwd())
@click.option('-o', '--output', type=click.File('w'))
def summarize(directory: str, output):
    """Summarize contents of training and evaluation."""
    summarize_results(directory, output)


if __name__ == '__main__':
    main()
