#! /usr/bin/env python
# coding=utf-8
# Copyright (c) 2019 Uber Technologies, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
import copy
import os
import random
import subprocess
import subprocess as sp
import sys
from collections import OrderedDict
from collections.abc import Mapping

import numpy

import ludwig.globals
from ludwig.utils.data_utils import figure_data_format


def get_experiment_description(
        config,
        dataset=None,
        training_set=None,
        validation_set=None,
        test_set=None,
        training_set_metadata=None,
        data_format=None,
        random_seed=None
):
    description = OrderedDict()
    description['ludwig_version'] = ludwig.globals.LUDWIG_VERSION
    description['command'] = ' '.join(sys.argv)

    try:
        with open(os.devnull, 'w') as devnull:
            is_a_git_repo = subprocess.call(['git', 'branch'],
                                            stderr=subprocess.STDOUT,
                                            stdout=devnull) == 0
        if is_a_git_repo:
            description['commit_hash'] = \
                subprocess.check_output(['git', 'rev-parse', 'HEAD']).decode(
                    'utf-8')[:12]
    except:
        pass

    if random_seed is not None:
        description['random_seed'] = random_seed

    if isinstance(dataset, str):
        description['dataset'] = dataset
    if isinstance(training_set, str):
        description['training_set'] = training_set
    if isinstance(validation_set, str):
        description['validation_set'] = validation_set
    if isinstance(test_set, str):
        description['test_set'] = test_set
    if training_set_metadata is not None:
        description['training_set_metadata'] = training_set_metadata

    # determine data format if not provided or auto
    if not data_format or data_format == 'auto':
        data_format = figure_data_format(
            dataset, training_set, validation_set, test_set
        )

    if data_format:
        description['data_format'] = str(data_format)

    description['config'] = config

    import tensorflow as tf
    description['tf_version'] = tf.__version__

    return description


def set_random_seed(random_seed):
    os.environ['PYTHONHASHSEED'] = str(random_seed)
    random.seed(random_seed)
    numpy.random.seed(random_seed)


def merge_dict(dct, merge_dct):
    """ Recursive dict merge. Inspired by :meth:``dict.update()``, instead of
    updating only top-level keys, dict_merge recurses down into dicts nested
    to an arbitrary depth, updating keys. The ``merge_dct`` is merged into
    ``dct``.
    :param dct: dict onto which the merge is executed
    :param merge_dct: dct merged into dct
    :return: None
    """
    dct = copy.deepcopy(dct)
    for k, v in merge_dct.items():
        if (k in dct and isinstance(dct[k], dict)
                and isinstance(merge_dct[k], Mapping)):
            dct[k] = merge_dict(dct[k], merge_dct[k])
        else:
            dct[k] = merge_dct[k]
    return dct


def sum_dicts(dicts, dict_type=dict):
    summed_dict = dict_type()
    for d in dicts:
        for key, value in d.items():
            if key in summed_dict:
                prev_value = summed_dict[key]
                if isinstance(value, (dict, OrderedDict)):
                    summed_dict[key] = sum_dicts([prev_value, value],
                                                 dict_type=type(value))
                elif isinstance(value, numpy.ndarray):
                    summed_dict[key] = numpy.concatenate((prev_value, value))
                else:
                    summed_dict[key] = prev_value + value
            else:
                summed_dict[key] = value
    return summed_dict


def resolve_pointers(dict1, dict2, dict2_name):
    resolved_dict = copy.deepcopy(dict1)
    for key in dict1:
        value = dict1[key]
        if value.startswith(dict2_name):
            key_in_dict2 = value[len(dict2_name):]
            if key_in_dict2 in dict2.keys():
                value = dict2[key_in_dict2]
                resolved_dict[key] = value
    return resolved_dict


def get_from_registry(key, registry):
    if hasattr(key, 'lower'):
        key = key.lower()
    if key in registry:
        return registry[key]
    else:
        raise ValueError(
            'Key {} not supported, available options: {}'.format(
                key, registry.keys()
            )
        )


def set_default_value(dictionary, key, value):
    if key not in dictionary:
        dictionary[key] = value


def set_default_values(dictionary, default_value_dictionary):
    # Set multiple default values
    for key, value in default_value_dictionary.items():
        set_default_value(dictionary, key, value)


def find_non_existing_dir_by_adding_suffix(directory_name):
    curr_directory_name = directory_name
    suffix = 0
    while os.path.exists(curr_directory_name):
        curr_directory_name = directory_name + '_' + str(suffix)
        suffix += 1
    return curr_directory_name


def get_class_attributes(c):
    return set(
        i for i in dir(c)
        if not callable(getattr(c, i)) and not i.startswith("_")
    )


def get_available_gpu_memory():
    _output_to_list = lambda x: x.decode('ascii').split('\n')[:-1]

    COMMAND = "nvidia-smi --query-gpu=memory.free --format=csv"
    try:
        memory_free_info = _output_to_list(sp.check_output(COMMAND.split()))[
                           1:]
        memory_free_values = [int(x.split()[0])
                              for i, x in enumerate(memory_free_info)]
    except Exception as e:
        print('"nvidia-smi" is probably not installed.', e)

    return memory_free_values


def get_output_directory(
        output_directory,
        experiment_name,
        model_name='run'
):
    base_dir_name = os.path.join(
        output_directory,
        experiment_name + ('_' if model_name else '') + model_name
    )
    return find_non_existing_dir_by_adding_suffix(base_dir_name)


def get_file_names(output_directory):
    description_fn = os.path.join(output_directory, 'description.json')
    training_stats_fn = os.path.join(
        output_directory, 'training_statistics.json')

    model_dir = os.path.join(output_directory, 'model')

    return description_fn, training_stats_fn, model_dir


def check_which_config(config, config_file):
    # check for config and config_file
    if config is None and config_file is None:
        raise ValueError(
            'Either config of config_file have to be'
            'not None to initialize a LudwigModel'
        )
    if config is not None and config_file is not None:
        raise ValueError(
            'Only one between config and '
            'config_file can be provided'
        )
    if not config:
        config = config_file
    return config
