import copyreg
import re
import warnings
from copy import deepcopy

import keras
import toolz as t
from keras import __version__ as keras_version
from keras import backend as K
from keras import optimizers
from keras.engine import topology
from keras.legacy import layers as legacy_layers
from keras.legacy import models as legacy_models
from keras.models import Sequential, model_from_config

import provenance as p
"""
Registers a pickling functions for Keras models so they can be used as inputs and
outputs to provenance'd functions. Currently the pickle functions do not account
for custom layers and functions (this will be possible to do though).

A custom value_repr is also added to make sure that keras models with identical
architecture and weights hash the same. If names are explicity given to layers in
the network and they are different then the hash will be different but if the
layers were anonymous care is taken to ensure that the hash is invariant to the
autogenerated names for these layers done by keras.
"""

assert keras_version.startswith('2'), 'provenance only supports keras 2'


# TODO: capture any custom_obects and return it in the dictionary
def reduce_model(model):
    """Returns a dict representing the state of the model.

    The dict contains:
        - the model's configuration (topology)
        - the model's weights
        - the model's optimizer's state (if any)

    Thus the model can be reinstantiated in
    the returned dict, without any of the code
    used for model definition or training.

    # Arguments
        model: Keras model instance.
    """

    def get_config(obj):
        # if obj is a serializable Keras class instance
        # e.g. optimizer, layer
        if hasattr(obj, 'get_config'):
            return {
                'class_name': obj.__class__.__name__,
                'config': obj.get_config()
            }
        return obj

    attrs = {}
    attrs['keras_version'] = str(keras_version).encode('utf8')
    attrs['backend'] = K.backend().encode('utf8')
    round_tripped_model = model_from_config(get_config(model))
    attrs['model_config'] = get_config(round_tripped_model)

    model_weights = {}
    attrs['model_weights'] = model_weights
    if legacy_models.needs_legacy_support(model):
        model_layers = legacy_models.legacy_sequential_layers(model)
    else:
        model_layers = model.layers

    model_weights[
        'layer_names'] = [layer.name.encode('utf8') for layer in model_layers]
    model_weights['backend'] = K.backend().encode('utf8')
    model_weights['keras_version'] = str(keras_version).encode('utf8')

    for layer in model_layers:
        g = {}
        model_weights[layer.name] = g
        symbolic_weights = layer.weights
        weight_values = K.batch_get_value(symbolic_weights)
        weight_names = []
        g['weight_names'] = weight_names
        for i, (w, val) in enumerate(zip(symbolic_weights, weight_values)):
            if hasattr(w, 'name') and w.name:
                name = str(w.name)
            else:
                name = 'param_' + str(i)
            weight_names.append(name)
            g[name] = val

    if hasattr(model, 'optimizer'):
        if isinstance(model.optimizer, optimizers.TFOptimizer):
            warnings.warn(
                'TensorFlow optimizers do not '
                'make it possible to access '
                'optimizer attributes or optimizer state '
                'after instantiation. '
                'As a result, we cannot save the optimizer '
                'as part of the model save file.'
                'You will have to compile your model again after loading it. '
                'Prefer using a Keras optimizer instead '
                '(see keras.io/optimizers).')
        else:
            attrs['training_config'] = t.valmap(get_config, {
                'optimizer_config':
                model.optimizer,
                'loss':
                model.loss,
                'metrics':
                model.metrics,
                'sample_weight_mode':
                model.sample_weight_mode,
                'loss_weights':
                model.loss_weights,
            })

            # Save optimizer weights.
            symbolic_weights = getattr(model.optimizer, 'weights')
            if symbolic_weights:
                optimizer_weights = {}
                attrs['optimizer_weigts'] = optimizer_weights
                weight_values = K.batch_get_value(symbolic_weights)
                for i, (w, val
                        ) in enumerate(zip(symbolic_weights, weight_values)):
                    # Default values of symbolic_weights is /variable for theano
                    if K.backend() == 'theano':
                        if hasattr(w, 'name') and w.name != "/variable":
                            name = str(w.name)
                        else:
                            name = 'param_' + str(i)
                    else:
                        if hasattr(w, 'name') and w.name:
                            name = str(w.name)
                        else:
                            name = 'param_' + str(i)
                    optimizer_weights[name] = val
    return attrs


def _load_weights(weights, layers):
    """Implements topological (order-based) weight loading.

    # Arguments
        weights: A dict of weights as returned by reduce_model
        layers: a list of target layers.

    # Raises
        ValueError: in case of mismatch between provided layers
            and weights file.
    """
    if 'keras_version' in weights:
        original_keras_version = weights['keras_version'].decode('utf8')
    else:
        original_keras_version = '1'
    if 'backend' in weights:
        original_backend = weights['backend'].decode('utf8')
    else:
        original_backend = None

    filtered_layers = []
    for layer in layers:
        lweights = layer.weights
        if lweights:
            filtered_layers.append(layer)

    layer_names = [n.decode('utf8') for n in weights['layer_names']]
    filtered_layer_names = []
    for name in layer_names:
        g = weights[name]
        weight_names = g['weight_names']
        if weight_names:
            filtered_layer_names.append(name)
    layer_names = filtered_layer_names
    if len(layer_names) != len(filtered_layers):
        raise ValueError('You are trying to load a weight file '
                         'containing ' + str(len(layer_names)) +
                         ' layers into a model with ' + str(
                             len(filtered_layers)) + ' layers.')

    # We batch weight value assignments in a single backend call
    # which provides a speedup in TensorFlow.
    weight_value_tuples = []
    for k, name in enumerate(layer_names):
        g = weights[name]
        weight_names = g['weight_names']
        weight_values = [g[weight_name] for weight_name in weight_names]
        layer = filtered_layers[k]
        symbolic_weights = layer.weights
        weight_values = topology.preprocess_weights_for_loading(
            layer, weight_values, original_keras_version, original_backend)
        if len(weight_values) != len(symbolic_weights):
            raise ValueError(
                'Layer #' + str(k) + ' (named "' + layer.name +
                '" in the current model) was found to '
                'correspond to layer ' + name + ' in the save file. '
                'However the new layer ' + layer.name + ' expects ' + str(
                    len(symbolic_weights)) +
                ' weights, but the saved weights have ' + str(
                    len(weight_values)) + ' elements.')
        weight_value_tuples += zip(symbolic_weights, weight_values)
    K.batch_set_value(weight_value_tuples)


REGISTERED_CUSTOM_OBJECTS = {}

def register_custom_objects(mapping, merge=False):
    if merge:
        res = t.merge(REGISTERED_CUSTOM_OBJECTS, mapping)
    else:
        res = mapping

    REGISTERED_CUSTOM_OBJECTS = res



#TODO: move custom_objects into the attrs
def load_model(attrs, custom_objects=None):
    """Loads a model saved via `reduce_model`.

    This was largely based on keras.models.load_model

    # Arguments
        attrs: Model attributes as given from reduce_model.
        custom_objects: Optional dictionary mapping names
            (strings) to custom classes or functions to be
            considered during deserialization.

    # Returns
        A Keras model instance. If an optimizer was found
        as part of the saved model, the model is already
        compiled. Otherwise, the model is uncompiled and
        a warning will be displayed.
    """
    if not custom_objects:
        custom_objects = {}

    custom_objects = t.merge(REGISTERED_CUSTOM_OBJECTS,
                             custom_objects)

    def convert_custom_objects(obj):
        """Handles custom object lookup.

        # Arguments
            obj: object, dict, or list.

        # Returns
            The same structure, where occurences
                of a custom object name have been replaced
                with the custom object.
        """
        if isinstance(obj, list):
            deserialized = []
            for value in obj:
                if value in custom_objects:
                    deserialized.append(custom_objects[value])
                else:
                    deserialized.append(value)
            return deserialized
        if isinstance(obj, dict):
            deserialized = {}
            for key, value in obj.items():
                if value in custom_objects:
                    deserialized[key] = custom_objects[value]
                else:
                    deserialized[key] = value
            return deserialized
        if obj in custom_objects:
            return custom_objects[obj]
        return obj

    # instantiate model
    model_config = attrs.get('model_config')
    if model_config is None:
        raise ValueError('No model found in config file.')
    model = model_from_config(model_config, custom_objects=custom_objects)

    # set weights
    _load_weights(attrs['model_weights'], model.layers)

    # instantiate optimizer
    training_config = attrs.get('training_config')
    if training_config is None:
        warnings.warn('No training configuration found in save file: '
                      'the model was *not* compiled. Compile it manually.')
        return model
    optimizer_config = training_config['optimizer_config']
    optimizer = optimizers.deserialize(
        optimizer_config, custom_objects=custom_objects)

    # Recover loss functions and metrics.
    loss = convert_custom_objects(training_config['loss'])
    metrics = convert_custom_objects(training_config['metrics'])
    sample_weight_mode = training_config['sample_weight_mode']
    loss_weights = training_config['loss_weights']

    # Compile model.
    model.compile(
        optimizer=optimizer,
        loss=loss,
        metrics=metrics,
        loss_weights=loss_weights,
        sample_weight_mode=sample_weight_mode)

    # Set optimizer weights.
    if 'optimizer_weights' in attrs:
        # Build train function (to get weight updates).
        if isinstance(model, keras.models.Sequential):
            model.model._make_train_function()
        else:
            model._make_train_function()
        optimizer_weights = attrs['optimizer_weights']
        model.optimizer.set_weights(optimizer_weight_values)
    return model


def copy_model(model):
    return load_model(reduce_model(model))


def pickle_model(model):
    return load_model, (reduce_model(model), )


copyreg.pickle(keras.engine.training.Model, pickle_model)
copyreg.pickle(Sequential, pickle_model)


def _update_layer_ref_name(layer_arrays):
    """
    e.g.
    the inbound_nodes in a Lambda is are layer_arrays
                        'inbound_nodes': [[['sequential_2', 1, 0, {}],
                                           ['sequential_2', 2, 0, {}]]],
    or the input_layers for a Model:
    'input_layers': [['input_AUTOGENERATED', 0, 0],
                     ['input_AUTOGENERATED', 0, 0]]

    """
    for layer in layer_arrays:
        layer_name = layer[0]
        if isinstance(layer_name, list):
            _update_layer_ref_name(layer)
        else:
            # This assumes that only autogenerated layers follow this pattern..
            # which may not hold 100% but seems harmless otherwise
            layer[0] = re.sub('_\d+$', '_AUTOGENERATED', layer_name)


def _layer_config_hash_repr(model_config, original_layer_names):
    class_name = topology._to_snake_case(model_config['class_name'])

    def record_layer_name(name):
        if len(original_layer_names) == 0 or (
                len(original_layer_names) > 0 and
                original_layer_names[-1] != name):
            original_layer_names.append(name)

    if class_name == 'input_layer':
        # the InputLayer has a custom prefix
        class_name = 'input'

    if model_config.get('name') and model_config['name'].startswith(
            class_name):
        record_layer_name(model_config['name'])
        model_config['name'] = class_name + "_AUTOGENERATED"

    config = model_config['config']

    if isinstance(config, dict):
        name = config.get('name', None)

        if name is not None:
            record_layer_name(name)

            if name.startswith(class_name):
                config['name'] = class_name + "_AUTOGENERATED"

    for layer_array_ref in ['input_layers', 'output_layers']:
        if layer_array_ref in config:
            _update_layer_ref_name(config[layer_array_ref])

    if 'layers' in config:  # e.g. Model
        for layer in config['layers']:
            _layer_config_hash_repr(layer, original_layer_names)

    if 'inbound_nodes' in model_config:
        _update_layer_ref_name(model_config['inbound_nodes'])

    if class_name == "sequential":
        for layer in config:
            _layer_config_hash_repr(layer, original_layer_names)


def model_config_hash_repr(model_config):
    model_config = deepcopy(model_config)
    original_layer_names = []
    _layer_config_hash_repr(model_config, original_layer_names)
    return model_config, original_layer_names


@p.value_repr.register(keras.engine.training.Model)
def model_hash_repr(model):
    attrs = reduce_model(model)
    # the keras version shouldn't change the hash typically since keras is good about backward compat
    del attrs['keras_version']

    # update the layers so the hash is invariant to autogenerated layer names
    model_config, original_layer_names = model_config_hash_repr(
        attrs['model_config'])
    attrs['model_config'] = model_config

    # we sort the weights by name to preserve order of the params (e.g. bias, kernel)
    weights = attrs['model_weights']
    if 'keras_version' in weights:
        del weights['keras_version']
    sanitized_weights = []
    for n in original_layer_names:
        if n not in weights:
            # model names are also returned in the original_layer_names
            # so we can safely skip them
            continue
        wd = weights[n]
        del wd['weight_names']
        sanitized_weights.append(
            [v for k, v in sorted(wd.items(), key=lambda i: i[0])])

    attrs['model_weights'] = sanitized_weights

    # do we need to do something similar for optimizer weights?
    return attrs
