# -*- coding: utf-8 -*-
import random
import numpy as np
from tensorflow import set_random_seed
import keras
from keras.models import Model
from keras.layers import *
from .attention import Attention

random.seed = 16
np.random.seed(16)
set_random_seed(16)


def precision(y_true, y_pred):
    """Precision metric.
    Only computes a batch-wise average of precision.
    Computes the precision, a metric for multi-label classification of
    how many selected items are relevant.
    """
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
    precision = true_positives / (predicted_positives + K.epsilon())
    return precision


def recall(y_true, y_pred):
    """Recall metric.
    Only computes a batch-wise average of recall.
    Computes the recall, a metric for multi-label classification of
    how many relevant items are selected.
    """
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
    recall = true_positives / (possible_positives + K.epsilon())
    return recall


def fbeta(y_true, y_pred, beta=1):
    """Computes the F score.
    The F score is the weighted harmonic mean of precision and recall.
    Here it is only computed as a batch-wise average, not globally.
    This is useful for multi-label classification, where input samples can be
    classified as sets of labels. By only using accuracy (precision) a model
    would achieve a perfect score by simply assigning every class to every
    input. In order to avoid this, a metric should penalize incorrect class
    assignments as well (recall). The F-beta score (ranged from 0.0 to 1.0)
    computes this, as a weighted mean of the proportion of correct class
    assignments vs. the proportion of incorrect class assignments.
    With beta = 1, this is equivalent to a F-measure. With beta < 1, assigning
    correct classes becomes more important, and with beta > 1 the metric is
    instead weighted towards penalizing incorrect class assignments.
    """
    if beta < 0:
        raise ValueError('The lowest choosable beta is zero (only precision).')

    p = precision(y_true, y_pred)
    r = recall(y_true, y_pred)
    bb = beta ** 2
    res = (1 + bb) * (p * r) / (bb * p + r + K.epsilon())
    return res


def softmax(x, axis=-1):
    ex = K.exp(x - K.max(x, axis=axis, keepdims=True))
    return ex / K.sum(ex, axis=axis, keepdims=True)


def squash(x, axis=-1):
    """From hinton's paper
    (||x||^2 / (1 + ||x||^2)) * (x / ||x||)
    """
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    scale = s_squared_norm / (1 + s_squared_norm)
    return scale * x / K.sqrt(s_squared_norm + K.epsilon())

def squash_half(x, axis=-1):
    """Use 0.5 in stead of 1 in hinton's paper.
    if 1, the norm of vector will be zoomed out.
    if 0.5, the norm will be zoomed in while original norm is less than 0.5
    and be zoomed out while original norm is greater than 0.5.
    (||x||^2 / (0.5 + ||x||^2)) * (x / ||x||)
    """
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True) + K.epsilon()
    scale = K.sqrt(s_squared_norm)/ (0.5 + s_squared_norm)
    return scale * x / K.sqrt(s_squared_norm + K.epsilon())

def squash_no_compression(x, axis=-1):
    """Use 0 in stead of 1 in hinton's paper without compression.
    x / ||x||
    """
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    return x / K.sqrt(s_squared_norm + K.epsilon())

def squash_1x(x, axis=-1):
    """x / (1 + ||x||)
    """
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    return x / (1 + K.sqrt(s_squared_norm + K.epsilon()))

def squash_e(x, axis=-1):
    """(1-e^(-||x||)) * (x / ||x||)
    """
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    scale = 1 - K.exp(-K.sqrt(s_squared_norm))
    return scale * x / K.sqrt(s_squared_norm + K.epsilon())

def squash_tanh(x, axis=-1):
    """tanh||x|| * (x / ||x||)
    """
    s_squared_norm = K.sum(K.square(x), axis, keepdims=True)
    scale = K.tanh(-K.sqrt(s_squared_norm))
    return scale * x / K.sqrt(s_squared_norm + K.epsilon())


class Capsule(Layer):
    """A Capsule Implement with Pure Keras
    """
    def __init__(self, num_capsule, dim_capsule, routings=3, kernel_size=(9, 1), share_weights=True,
                 activation='default', **kwargs):
        super(Capsule, self).__init__(**kwargs)
        self.num_capsule = num_capsule
        self.dim_capsule = dim_capsule
        self.routings = routings
        self.kernel_size = kernel_size
        self.share_weights = share_weights
        if activation == 'default':
            self.activation = squash_no_compression
        else:
            self.activation = Activation(activation)

    def build(self, input_shape):
        super(Capsule, self).build(input_shape)
        input_dim_capsule = input_shape[-1]
        if self.share_weights:
            self.W = self.add_weight(name='capsule_kernel',
                                     shape=(1, input_dim_capsule,
                                            self.num_capsule * self.dim_capsule),
                                     # shape=self.kernel_size,
                                     initializer='glorot_uniform',
                                     trainable=True)
        else:
            input_num_capsule = input_shape[-2]
            self.W = self.add_weight(name='capsule_kernel',
                                     shape=(input_num_capsule,
                                            input_dim_capsule,
                                            self.num_capsule * self.dim_capsule),
                                     initializer='glorot_uniform',
                                     trainable=True)

    def call(self, u_vecs):
        if self.share_weights:
            u_hat_vecs = K.conv1d(u_vecs, self.W)
        else:
            u_hat_vecs = K.local_conv1d(u_vecs, self.W, [1], [1])

        batch_size = K.shape(u_vecs)[0]
        input_num_capsule = K.shape(u_vecs)[1]
        u_hat_vecs = K.reshape(u_hat_vecs, (batch_size, input_num_capsule,
                                            self.num_capsule, self.dim_capsule))
        u_hat_vecs = K.permute_dimensions(u_hat_vecs, (0, 2, 1, 3))
        # final u_hat_vecs.shape = [None, num_capsule, input_num_capsule, dim_capsule]

        b = K.zeros_like(u_hat_vecs[:, :, :, 0])  # shape = [None, num_capsule, input_num_capsule]
        for i in range(self.routings):
            b = K.permute_dimensions(b, (0, 2, 1))  # shape = [None, input_num_capsule, num_capsule]
            c = K.softmax(b)
            c = K.permute_dimensions(c, (0, 2, 1))
            b = K.permute_dimensions(b, (0, 2, 1))
            outputs = self.activation(K.batch_dot(c, u_hat_vecs, [2, 2]))
            if i < self.routings - 1:
                b = K.batch_dot(outputs, u_hat_vecs, [2, 3])

        return outputs

    def compute_output_shape(self, input_shape):
        return (None, self.num_capsule, self.dim_capsule)


class TextClassifier():

    def model(self, embeddings_matrix, maxlen, word_index, num_class):
        input1 = Input(shape=(maxlen,))
        embed_layer = Embedding(len(word_index) + 1,
                                embeddings_matrix.shape[1],
                                input_length=maxlen,
                                weights=[embeddings_matrix],
                                trainable=True)(input1)
        embed_layer = SpatialDropout1D(0.2)(embed_layer) # 0.28

        x = Bidirectional(CuDNNGRU(128, return_sequences=True))(embed_layer)
        x = Activation('relu')(x)
        x = Dropout(0.2)(x) # 0.25
        x = Bidirectional(CuDNNGRU(128,  return_sequences=True))(x)
        x = Activation('relu')(x)
        x = Dropout(0.2)(x) # 0.25
        capsule = Capsule(num_capsule=10, dim_capsule=16, routings=5,
                          share_weights=True)(x)
        capsule = Flatten()(capsule)
        capsule = Dropout(0.2)(capsule) # 0.25
        if num_class == 2:
            output = Dense(num_class, activation='sigmoid')(capsule)
            loss = 'binary_crossentropy'
        else:
            output = Dense(num_class, activation='softmax')(capsule)
            loss = 'categorical_crossentropy'
        model = Model(inputs=input1, outputs=output)
        model.compile(
            loss=loss,
            optimizer='adam',
            metrics=["categorical_accuracy"])
        return model