import copy

import numpy as np
import tensorflow.keras as tfk
import tensorflow.keras.layers as tfkl
from tensorflow.python.keras import activations
from tensorflow.python.keras import constraints
from tensorflow.python.keras import initializers
from tensorflow.python.ops import nn

from dynastes import regularizers
from dynastes import weight_normalizers


class _WscaleInitializer(tfk.initializers.Initializer):

    def __init__(self,
                 initializer,
                 lrmul=1.,
                 **kwargs):
        super(_WscaleInitializer, self).__init__(**kwargs)
        self.initializer = tfk.initializers.get(initializer)
        self.lrmul = lrmul

    def __call__(self, shape, dtype=None):
        return self.initializer(shape, dtype=dtype) / self.lrmul

    def get_config(self):
        config = {
            'initializer': tfk.initializers.serialize(self.initializer),
            'lrmul': self.lrmul,
        }
        base_config = super(_WscaleInitializer, self).get_config()
        return {**base_config, **config}


def _get_regularizers_from_keywords(kwargs):
    _initializers = {}
    _regularizers = {}
    _constraints = {}
    _normalizers = {}

    kwarg_keys = copy.copy(list(kwargs.keys()))

    for kwarg in kwarg_keys:
        if kwarg.endswith('initializer'):
            _initializers[kwarg.split('_initializer')[0]] = initializers.get(kwargs.pop(kwarg, None))
        elif kwarg.endswith('regularizer'):
            if kwarg != 'activity_regularizer':
                _regularizers[kwarg.split('_regularizer')[0]] = regularizers.get(kwargs.pop(kwarg, None))
        elif kwarg.endswith('constraint'):
            _constraints[kwarg.split('_constraint')[0]] = constraints.get(kwargs.pop(kwarg, None))
        elif kwarg.endswith('normalizer'):
            _normalizers[kwarg.split('_normalizer')[0]] = weight_normalizers.get(kwargs.pop(kwarg, None))

    return _initializers, _regularizers, _constraints, _normalizers


class DynastesBaseLayer(tfkl.Layer):
    def __init__(self,
                 activity_regularizer=None,
                 trainable=True,
                 name=None,
                 use_wscale=False,
                 wlrmul=1.,
                 wgain=np.sqrt(2),
                 **kwargs):
        self.weights_dict = {}
        self.initializers, self.regularizers, self.constraints, self.normalizers = _get_regularizers_from_keywords(
            kwargs)
        self.use_wscale = use_wscale
        self.lrmul = wlrmul
        self.gain = wgain
        super(DynastesBaseLayer, self).__init__(
            trainable=trainable,
            name=name,
            activity_regularizer=regularizers.get(activity_regularizer),
            **kwargs)

    def get_initializer(self, name):
        if name not in self.initializers:
            if name == 'kernel':
                self.initializers['kernel'] = initializers.get('he_uniform')
            elif name == 'bias':
                self.initializers['bias'] = initializers.get('zeros')
            else:
                self.initializers[name] = initializers.get('glorot_uniform')

        return self.initializers[name]

    def get_regularizer(self, name):
        if name not in self.regularizers:
            self.regularizers[name] = regularizers.get(None)
        return self.regularizers[name]

    def get_constraint(self, name):
        if name not in self.constraints:
            self.constraints[name] = constraints.get(None)
        return self.constraints[name]

    def add_weight(self,
                   name=None,
                   shape=None,
                   trainable=None,
                   partitioner=None,
                   initializer=None,
                   regularizer=None,
                   constraint=None,
                   dtype=None,
                   use_resource=None,
                   **kwargs):
        if initializer is not None:
            self.initializers[name] = initializers.get(initializer)
        if regularizer is not None:
            self.regularizers[name] = regularizers.get(regularizer)
        if constraint is not None:
            self.constraints[name] = constraints.get(constraint)
        _initializer = self.get_initializer(name)
        if self.use_wscale:
            _initializer = _WscaleInitializer(_initializer, lrmul=self.lrmul)
            self.initializers[name] = _initializer
            if name in self.normalizers and self.normalizers[name] is not None:
                self.normalizers[name] = weight_normalizers.WscaleNormalizer(next_layer=self.normalizers[name], lrmul=self.lrmul,
                                                          gain=self.gain)
            else:
                self.normalizers[name] = weight_normalizers.WscaleNormalizer(lrmul=self.lrmul, gain=self.gain)

        weight = super(DynastesBaseLayer, self).add_weight(name=name,
                                                           shape=shape,
                                                           initializer=_initializer,
                                                           regularizer=self.get_regularizer(name),
                                                           trainable=trainable,
                                                           constraint=self.get_constraint(name),
                                                           partitioner=partitioner,
                                                           use_resource=use_resource,
                                                           **kwargs)
        if name in self.normalizers:
            if self.normalizers[name] is not None:
                self.normalizers[name].build(shape)
        self.weights_dict[name] = weight
        return weight

    def get_weight(self, name, training=None):
        w = self.weights_dict[name]
        if name in self.normalizers and self.normalizers[name] is not None:
            w = self.normalizers[name](w, training=training)
        return w

    def get_config(self):
        config = {
            'activity_regularizer': regularizers.serialize(self.activity_regularizer),
        }
        for name, initializer in self.initializers.items():
            config[name + '_initializer'] = initializers.serialize(initializer)
        for name, regularizer in self.regularizers.items():
            config[name + '_regularizer'] = regularizers.serialize(regularizer)
        for name, constraint in self.constraints.items():
            config[name + '_constraint'] = constraints.serialize(constraint)
        for name, normalizer in self.normalizers.items():
            config[name + '_normalizer'] = weight_normalizers.serialize(normalizer)

        base_config = super(DynastesBaseLayer, self).get_config()
        return {**base_config, **config}


class ActivatedKernelBiasBaseLayer(DynastesBaseLayer):
    def __init__(self,
                 activation=None,
                 use_bias=True,
                 activity_regularizer=None,
                 **kwargs):
        super(ActivatedKernelBiasBaseLayer, self).__init__(
            activity_regularizer=regularizers.get(activity_regularizer),
            **kwargs)
        self.activation = activations.get(activation)
        self.use_bias = use_bias

    def build_kernel(self, shape):
        return self.add_weight(
            name='kernel',
            shape=shape,
            trainable=True,
            dtype=self.dtype)

    def build_bias(self, output_dim):
        if self.use_bias:
            return self.add_weight(
                name='bias',
                shape=(output_dim,),
                trainable=True,
                dtype=self.dtype)
        return None

    def post_process_call(self, x, training=None):
        if self.use_bias:
            x = nn.bias_add(x, self.get_weight('bias', training=training), data_format='NHWC')
        if self.activation is not None:
            return self.activation(x)
        return x

    def call(self, x, training=None):
        return self.post_process_call(x, training=training)

    def get_config(self):
        config = {
            'activation': activations.serialize(self.activation),
            'use_bias': self.use_bias,
        }
        base_config = super(ActivatedKernelBiasBaseLayer, self).get_config()
        return {**base_config, **config}
