import json

import tensorflow as tf
from tensorflow import keras
import tensorflow.keras.backend as K

from keras_bert.keras_pos_embd import PositionEmbedding
from keras_bert.layers import get_inputs, get_embedding, TokenEmbedding, EmbeddingSimilarity, Masked, Extract
from keras_bert.keras_layer_normalization import LayerNormalization
from keras_bert.keras_multi_head import MultiHeadAttention
from keras_bert.keras_position_wise_feed_forward import FeedForward

def gelu(x):
    return 0.5 * x * (1.0 + tf.math.erf(x / tf.sqrt(2.0)))

class Bert(keras.Model):
    def __init__(
            self,
            token_num,
            pos_num=512,
            seq_len=512,
            embed_dim=768,
            transformer_num=12,
            head_num=12,
            feed_forward_dim=3072,
            dropout_rate=0.1,
            attention_activation=None,
            feed_forward_activation=gelu,
            custom_layers=None,
            training=True,
            trainable=None,
            lr=1e-4,
            name='Bert'):
        super().__init__(name=name)
        self.token_num = token_num
        self.pos_num = pos_num
        self.seq_len = seq_len
        self.embed_dim = embed_dim
        self.transformer_num = transformer_num
        self.head_num = head_num
        self.feed_forward_dim = feed_forward_dim
        self.dropout_rate = dropout_rate
        self.attention_activation = attention_activation
        self.feed_forward_activation = feed_forward_activation
        self.custom_layers = custom_layers
        self.training = training
        self.trainable = trainable
        self.lr = lr

        # build layers
        # embedding
        self.token_embedding_layer = TokenEmbedding(
            input_dim=token_num,
            output_dim=embed_dim,
            mask_zero=True,
            trainable=trainable,
            name='Embedding-Token',
        )
        self.segment_embedding_layer = keras.layers.Embedding(
            input_dim=2,
            output_dim=embed_dim,
            trainable=trainable,
            name='Embedding-Segment',
        )
        self.position_embedding_layer = PositionEmbedding(
            input_dim=pos_num,
            output_dim=embed_dim,
            mode=PositionEmbedding.MODE_ADD,
            trainable=trainable,
            name='Embedding-Position',
        )
        self.embedding_layer_norm = LayerNormalization(
            trainable=trainable,
            name='Embedding-Norm',
        )

        self.encoder_multihead_layers = []
        self.encoder_ffn_layers = []
        self.encoder_attention_norm = []
        self.encoder_ffn_norm = []
        # attention layers
        for i in range(transformer_num):
            base_name = 'Encoder-%d' % (i + 1)
            attention_name = '%s-MultiHeadSelfAttention' % base_name
            feed_forward_name = '%s-FeedForward' % base_name
            self.encoder_multihead_layers.append(MultiHeadAttention(
                head_num=head_num,
                activation=attention_activation,
                history_only=False,
                trainable=trainable,
                name=attention_name,
            ))
            self.encoder_ffn_layers.append(FeedForward(
                units=feed_forward_dim,
                activation=feed_forward_activation,
                trainable=trainable,
                name=feed_forward_name,
            ))
            self.encoder_attention_norm.append(LayerNormalization(
                trainable=trainable,
                name='%s-Norm' % attention_name,
            ))
            self.encoder_ffn_norm.append(LayerNormalization(
                trainable=trainable,
                name='%s-Norm' % feed_forward_name,
            ))

    def call(self, inputs):

        embeddings = [
            self.token_embedding_layer(inputs[0]),
            self.segment_embedding_layer(inputs[1])
        ]
        embeddings[0], embed_weights = embeddings[0]
        embed_layer = keras.layers.Add(
            name='Embedding-Token-Segment')(embeddings)
        embed_layer = self.position_embedding_layer(embed_layer)

        if self.dropout_rate > 0.0:
            dropout_layer = keras.layers.Dropout(
                rate=self.dropout_rate,
                name='Embedding-Dropout',
            )(embed_layer)
        else:
            dropout_layer = embed_layer

        embedding_output = self.embedding_layer_norm(dropout_layer)

        def _wrap_layer(name, input_layer, build_func, norm_layer, dropout_rate=0.0, trainable=True):
            """Wrap layers with residual, normalization and dropout.

            :param name: Prefix of names for internal layers.
            :param input_layer: Input layer.
            :param build_func: A callable that takes the input tensor and generates the output tensor.
            :param dropout_rate: Dropout rate.
            :param trainable: Whether the layers are trainable.
            :return: Output layer.
            """
            build_output = build_func(input_layer)
            if dropout_rate > 0.0:
                dropout_layer = keras.layers.Dropout(
                    rate=dropout_rate,
                    name='%s-Dropout' % name,
                )(build_output)
            else:
                dropout_layer = build_output
            if isinstance(input_layer, list):
                input_layer = input_layer[0]
            add_layer = keras.layers.Add(name='%s-Add' %
                                         name)([input_layer, dropout_layer])
            normal_layer = norm_layer(add_layer)
            return normal_layer

        last_layer = embedding_output
        output_tensor_list = [last_layer]
        # self attention
        for i in range(self.transformer_num):
            base_name = 'Encoder-%d' % (i + 1)
            attention_name = '%s-MultiHeadSelfAttention' % base_name
            feed_forward_name = '%s-FeedForward' % base_name
            self_attention_output = _wrap_layer(
                name=attention_name,
                input_layer=last_layer,
                build_func=self.encoder_multihead_layers[i],
                norm_layer=self.encoder_attention_norm[i],
                dropout_rate=self.dropout_rate,
                trainable=self.trainable)
            last_layer = _wrap_layer(
                name=attention_name,
                input_layer=self_attention_output,
                build_func=self.encoder_ffn_layers[i],
                norm_layer=self.encoder_ffn_norm[i],
                dropout_rate=self.dropout_rate,
                trainable=self.trainable)
            output_tensor_list.append(last_layer)

        return output_tensor_list


def build_model_from_config(config_file,
                            training=False,
                            trainable=None,
                            seq_len=None,
                            build=True):
    """Build the model from config file.
    :param config_file: The path to the JSON configuration file.
    :param training: If training, the whole model will be returned.
    :param trainable: Whether the model is trainable.
    :param seq_len: If it is not None and it is shorter than the value in the config file, the weights in
                    position embeddings will be sliced to fit the new length.
    :return: model and config
    """
    with open(config_file, 'r') as reader:
        config = json.loads(reader.read())
    if seq_len is not None:
        config['max_position_embeddings'] = min(
            seq_len, config['max_position_embeddings'])
    if trainable is None:
        trainable = training
    model = Bert(
        token_num=config['vocab_size'],
        pos_num=config['max_position_embeddings'],
        seq_len=config['max_position_embeddings'],
        embed_dim=config['hidden_size'],
        transformer_num=config['num_hidden_layers'],
        head_num=config['num_attention_heads'],
        feed_forward_dim=config['intermediate_size'],
        training=training,
        trainable=trainable,
    )
    if build:
        model.build(input_shape=[(None, None), (None, None), (None, None)])
    return model, config
