from keras import backend as K
from keras import initializations
from keras.backend.common import _EPSILON 	
from keras.engine.topology import Layer
from keras.engine import InputSpec

from theano.tensor.nnet import h_softmax
import theano.tensor as T



class HierarchicalSoftmax(Layer): 

    def __init__(self, output_dim, init='glorot_uniform', **kwargs):
        self.init = initializations.get(init)
        self.output_dim = output_dim
        
        def hshape(n):
            from math import sqrt, ceil
            l1 = ceil(sqrt(n))
            l2 = ceil(n / l1)
            return int(l1), int(l2)
      
        self.n_classes, self.n_outputs_per_class = hshape(output_dim)
        input_dim =  self.input_spec[0].shape[-1]
        self.W1 = self.init((input_dim, self.n_classes), name='{}_W1'.format(self.name))
        self.b1 = K.zeros((self.n_classes,),  name='{}_b1'.format(self.name))
        self.W2 = self.init((self.n_classes, input_dim, self.n_outputs_per_class), name='{}_W2'.format(self.name))
        self.b2 = K.zeros((self.n_classes, self.n_outputs_per_class),  name='{}_b2'.format(self.name))

        self.trainable_weights = [self.W1, self.b1, self.W2, self.b2]


    
    def get_output_shape_for(self, input_shape):
        return (input_shape[0][0], input_shape[0][1], None)

    def call(self, X, mask=None):
       
        input_shape = self.input_spec[0].shape
        
        x = K.reshape(X[0], (-1, input_shape[2]))
        target = X[1].flatten() if self.trainable else None
        
        Y = h_softmax(x, K.shape(x)[0], self.output_dim, 
                              self.n_classes, self.n_outputs_per_class,
                              self.W1, self.b1, self.W2, self.b2, target)
        
        output_dim = 1 if self.trainable else self.output_dim    
        input_length = K.shape(X[0])[1]
       
        y = K.reshape(Y, (-1, input_length, output_dim))
        return y
        
    

    def get_config(self):
        config = {'output_dim': self.output_dim,
                  'init': self.init.__name__}
        base_config = super(HierarchicalSoftmax, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))
        
def hs_categorical_crossentropy(y_true, y_pred):
    y_pred = T.clip(y_pred, _EPSILON, 1.0 - _EPSILON)    
    return T.nnet.categorical_crossentropy(y_pred, y_true)
        super(HierarchicalSoftmax, self).__init__(**kwargs)

    def build(self, input_shape):
        
        self.input_spec = [InputSpec(shape=shape) for shape in input_shape]