"""
Generate graph embedding based on Pytorch BigGraph library  

"""

import os
import sys
import typing
from pathlib import Path
from argparse import Namespace
from kgtk.cli_argparse import KGTKArgumentParser, KGTKFiles

# remove the Issue: Initializing libiomp5.dylib, but found libiomp5.dylib already initialized.
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'


def parser():
    return {
        'help': 'graph embedding functionality',
        'description': 'Generate graph embedding in kgtk tsv format, here we use PytorchBigGraph as \
        low-level implementation '
    }


def add_arguments_extended(parser: KGTKArgumentParser, parsed_shared_args: Namespace):
    """
    Parse arguments
    Args:
        parser (argparse.ArgumentParser)
    """
    # import modules locally
    from kgtk.io.kgtkreader import KgtkReader, KgtkReaderOptions, KgtkReaderMode
    from kgtk.utils.argparsehelpers import optional_bool
    from kgtk.value.kgtkvalueoptions import KgtkValueOptions

    _expert: bool = parsed_shared_args._expert

    # IO
    parser.add_input_file()
    parser.add_output_file()

    parser.add_argument('-l', "--log", dest="log_file_path",
                        help="Setting the log path [Default: None]",
                        type=Path, default=None, metavar="")
    parser.add_argument('-T', '--temporary_directory', dest='temporary_directory',
                        help="Sepecify the directory location to store temporary file",
                        type=Path, default=Path('/tmp/'), metavar='')
    parser.add_argument('-ot', '--output_format', dest='output_format',
                        help="Outputformat for embeddings [Default: w2v] Choice: kgtk | w2v | glove",
                        default='w2v', metavar='')
    parser.add_argument('-r', '--retain_temporary_data', dest='retain_temporary_data',
                        help="When opearte graph, some tempory files will be generated, set True to retain these files",
                        type=bool, default=True, metavar='True|False')
    # Training parameters
    parser.add_argument('-d', "--dimension", dest="dimension_num",
                        help="Dimension of the real space the embedding live in [Default: 100]",
                        type=int, default=100, metavar="")
    parser.add_argument('-s', "--init_scale", dest="init_scale",
                        help="Generating the initial embedding with this standard deviation [Default: 0.001]" +
                             "If no initial embeddings are provided, they are generated by sampling each dimension" +
                             "from a centered normal distribution having this standard deviation.",
                        type=float, default=0.001, metavar="")
    parser.add_argument('-c', '--comparator', dest='comparator',
                        help="How the embeddings of the two sides of an edge (after having already " +
                             "undergone some processing) are compared to each other to produce a score[Default: dot]," +
                             "Choice: dot|cos|l2|squared_l2",
                        default='dot', choices=['dot', 'cos', 'l2', 'squared_l2'], metavar='dot|cos|l2|squared_l2')
    parser.add_argument('-op', '--operator', dest='operator',
                        help="The transformation to apply to the embedding of one of the sides of the edge " +
                             "(typically the right-hand one) before comparing it with the other one. It reflects"
                             "which model that embedding uses. [Default:ComplEx]", default='ComplEx',
                        metavar='RESCAL|DistMult|ComplEx|TransE')
    parser.add_argument('-e', '--num_epochs', dest='num_epochs',
                        help="The number of times the training loop iterates over all the edges.[Default:100]",
                        type=int, default=100, metavar='')
    parser.add_argument('-b', '--bias', dest='bias',
                        help="Whether use the bias choice [Default: False],If enabled, withhold the first " +
                             "dimension of the embeddings from the comparator and instead use it as a bias, adding " +
                             "back to the score. Makes sense for logistic and softmax loss functions. ",
                        type=bool, default=False, metavar='True|False')
    parser.add_argument('-w', '--workers', dest='workers',
                        help="The number of worker processes for training. If not given, set to CPU count.",
                        type=int, default=None, metavar='')
    parser.add_argument('-bs', '--batch_size', dest='batch_size',
                        help="The number of edges per batch.[Default:1000]",
                        type=int, default=1000, metavar='')
    parser.add_argument('-lf', '--loss_fn', dest='loss_fn',
                        help="How the scores of positive edges and their corresponding negatives " +
                             "are evaluated.[Default: ranking], Choice: ranking|logistic|softmax",
                        # default will be setting to ranking later
                        default=None, choices=['ranking', 'logistic', 'softmax', None],
                        metavar='ranking|logistic|softmax')
    parser.add_argument('-lr', '--learning_rate', dest='learning_rate',
                        help="The learning rate for the optimizer.[Default: 0.1]",
                        # default will be setting to 0.1 later
                        type=float, default=None, metavar='')
    parser.add_argument('-ef', '--eval_fraction', dest='eval_fraction',
                        help="The fraction of edges withheld from training and used to track evaluation " +
                             "metrics during training. [Defalut:0.0 training all edges ]",
                        type=float, default=0.0, metavar='')
    parser.add_argument('-dr', '--dynamic_relations', dest='dynamic_relations',
                        help="Whether use dynamic relations (when graphs with a " +
                             "large number of relations) [Default: True]",
                        type=bool, default=True, metavar='True|False')
    parser.add_argument('-ge', '--global_emb', dest='global_emb',
                        help="Whether use global embedding, if enabled, add to each embedding a vector that is common "
                             "to all the entities of a certain type. This vector is learned "
                             "during training.[Default: False] ", type=bool, default=False, metavar='True|False')
    # kgtk format
    parser.add_argument("--no-output-header", dest="output_no_header", metavar="True|False",
                        help="When true, do not write a header to the output file (default=%(default)s).",
                        type=optional_bool, nargs='?', const=True, default=False)

    KgtkReader.add_debug_arguments(parser)
    KgtkReaderOptions.add_arguments(parser,
                                    mode_options=True,
                                    default_mode=KgtkReaderMode[parsed_shared_args._mode],
                                    expert=_expert)
    KgtkValueOptions.add_arguments(parser)


def run(input_file: KGTKFiles,
        output_file: KGTKFiles,
        **kwargs):
    """
    **kwargs stores all parameters providing by user
    """
    from kgtk.graph_embeddings.graph_embeddings import ComputeGraphEmbeddings
    from kgtk.io.kgtkreader import KgtkReaderOptions
    from kgtk.value.kgtkvalueoptions import KgtkValueOptions

    input_kgtk_file: Path = KGTKArgumentParser.get_input_file(input_file)
    output_kgtk_file: Path = KGTKArgumentParser.get_output_file(output_file)
    reader_options: KgtkReaderOptions = KgtkReaderOptions.from_dict(kwargs)
    value_options: KgtkValueOptions = KgtkValueOptions.from_dict(kwargs)
    error_file: typing.TextIO = sys.stdout if kwargs.get("errors_to_stdout") else sys.stderr

    temporary_directory = kwargs['temporary_directory']
    output_format = kwargs['output_format']
    retain_temporary_data = kwargs['retain_temporary_data']
    verbose = kwargs['verbose']
    very_verbose = kwargs['very_verbose']
    log_file_path = kwargs['log_file_path']

    operator = kwargs['operator']
    dynamic_relations = kwargs['dynamic_relations']
    dimension_num = kwargs['dimension_num']
    global_emb = kwargs['global_emb']
    comparator = kwargs['comparator']
    init_scale = kwargs['init_scale']
    bias = kwargs['bias']
    num_epochs = kwargs['num_epochs']
    loss_fn = kwargs['loss_fn']
    learning_rate = kwargs['learning_rate']
    eval_fraction = kwargs['eval_fraction']
    output_no_header = kwargs['output_no_header']

    cge = ComputeGraphEmbeddings(input_kgtk_file=input_kgtk_file,
                                 output_kgtk_file=output_kgtk_file,
                                 temporary_directory=temporary_directory,
                                 output_format=output_format,
                                 retain_temporary_data=retain_temporary_data,
                                 verbose=verbose,
                                 very_verbose=very_verbose,
                                 log_file_path=log_file_path,
                                 reader_options=reader_options,
                                 value_options=value_options,
                                 error_file=error_file,
                                 operator=operator,
                                 dynamic_relations=dynamic_relations,
                                 dimension_num=dimension_num,
                                 global_emb=global_emb,
                                 comparator=comparator,
                                 init_scale=init_scale,
                                 bias=bias,
                                 num_epochs=num_epochs,
                                 loss_fn=loss_fn,
                                 learning_rate=learning_rate,
                                 eval_fraction=eval_fraction,
                                 output_no_header=output_no_header)

    cge.process()
