# 
# Warning: GPU performance is worst than CPU performance
# I guess that is for reading weights from the tensor.
#
# BoschNNModel: Keras NN model based on Sequential Model
#    Stores input data from further processing
#    Stores weights for first and last layers 
#    Note: Use batch_size = 1 and shuffle = False. Do the shuffling before.
#    I'm not sure though.
#    
# BoschNNCallback: Callback for weight processing
#
# BoschNNInput: Custom layer based on Dense Layer. This should be the 'first' layer in BoschNNModel.
#    Handles np.NaN input values. Basically, set them to 0.  
#
# BoschNNOutput: Custom layer based on Dense Layer. This should be the 'last' layer in BoschNNModel.
#    I am not sure why this layer exists.
#


from keras import backend as K
from keras.engine.topology import Layer
import numpy as np

from keras.layers import Input, Dense, Activation, Conv1D
from keras.models import Sequential, Model
from keras.optimizers import RMSprop, SGD
from keras.callbacks import Callback, ProgbarLogger

import theano.tensor as T

#
# from keras.engine.training import _make_batches
#
def _make_batches(size, batch_size):
    num_batches = int(np.ceil(size / float(batch_size)))
    return [(i * batch_size, min(size, (i + 1) * batch_size))
            for i in range(0, num_batches)]
    
class BoschNNModel(Sequential):
    def __init__(self, layers=None, name=None, **kwargs):
        self.nanMask = None
        self.notNanMask = None
        self.X = -1
        self.y = -1
        self.current_batch = -1
        self.null_indices = -1
        self.samples = -1
        self.batch_size = -1 
        self.batches = -1
        self.input_w = -1
        self.output_w = -1
        super(BoschNNModel, self).__init__(layers, name, **kwargs)
        
    def fit(self, X, y, **kwargs):
# For proper training 'batch_size' == 1
# and 'shuffle' == False.        
#         assert (kwargs['batch_size']==1)
        assert (kwargs['shuffle']==False)
        self.X = X
        self.y = y
        super(BoschNNModel, self).fit(X, y, **kwargs)
        

class BoschNNCallback(Callback):
    def on_train_begin(self, logs={}):
#         print "Train Begin:"
        self.model.samples = self.params['samples']
        self.model.batch_size = self.params['batch_size']
        self.model.batches = _make_batches(self.model.samples, self.model.batch_size)
        return
    
    def on_train_end(self, logs={}):
        return

    def on_epoch_begin(self, epoch, logs={}):
        return

    def on_epoch_end(self, epoch, logs={}):
        return
    
    def on_batch_begin(self, batch, logs={}):
#         print "Batch Begin:"
        start_index = batch*self.model.batch_size 
        stop_index = min( (batch+1) * self.model.batch_size, self.model.samples)
        self.model.current_batch = (self.model.X[start_index:stop_index])[0]
        self.model.null_indices = np.nonzero(np.isnan(self.model.current_batch))
        self.model.input_w = self.model.get_layer(name='first').get_weights()[0]
        self.model.output_w = self.model.get_layer(name='last').get_weights()[0]
        return

    def on_batch_end(self, batch, logs={}):
# If we have any np.NaN in the batch.        
        if (len(self.model.null_indices[0]) > 0):
            input_w_new = self.model.get_layer(name='first').get_weights()[0]
            output_w_new = self.model.get_layer(name='last').get_weights()[0]
# I am supprized that for np.NaN input neurons, backpropagation doesn't change neuron weights.
# See below             
#             diff = self.model.input_w[self.model.null_indices] - input_w_new[self.model.null_indices]   
        return


class BoschNNInput(Dense):
#     def __init__(self, parentModel, units, **kwargs):
#         super(BoschNNInput, self).__init__(units=units, **kwargs)
#         self.parentModel = parentModel
    def __init__(self, **kwargs):
        super(BoschNNInput, self).__init__(**kwargs)
#         self.units = units
        
    def build(self, input_shape):
        super(BoschNNInput, self).build(input_shape)

    def call(self, inputs):
        nanMask = T.isnan(inputs)
        notNanMask =T.invert(nanMask)
# any special use for this?
#         self.parentModel.nanMask = nanMask
#         self.parentModel.nonNanMask = notNanMask
        
#         inputs_copy = inputs.copy()
#         T.set_subtensor(inputs_copy[nanMask.nonzero(),:], 0*inputs_copy[nanMask.nonzero(),:])
        inputs_copy = T.set_subtensor(inputs[nanMask.nonzero()], 0)
        return super(BoschNNInput, self).call(inputs_copy) 
#
#    Or else:
#         output = K.dot(inputs_copy, self.kernel)
#         if self.use_bias:
#             output = K.bias_add(output, self.bias)
#         if self.activation is not None:
#             output = self.activation(output)
#         return output
        
    def compute_output_shape(self, input_shape):
        return super(BoschNNInput, self).compute_output_shape(input_shape)

    def get_config(self):
        base_config = super(BoschNNInput, self).get_config()
#         base_config["units"] = self.units
#         base_config["parentModel"] = None
        return base_config
    
class BoschNNOutput(Dense):
    def __init__(self, parentModel, units, **kwargs):
        super(BoschNNOutput, self).__init__(units=units, **kwargs)
        self.parentModel = parentModel

    def build(self, input_shape):
        super(BoschNNOutput, self).build(input_shape)

#     def update_fn(self, x):
#         mask = self.parentModel.nanMask
#         maskinv = self.parentModel.notNanMask
#         
#         T.set_subtensor(self.kernel[mask.nonzero(),:], c*self.kernel[mask.nonzero(),:])
#         T.set_subtensor(self.kernel[mask_inv.nonzero(),:], 5*self.kernel[mask_inv.nonzero(),:])
#         
#         return self.kernel
    
    def call(self, inputs):
        return super(BoschNNOutput, self).call(inputs)

    def compute_output_shape(self, input_shape):
        return super(BoschNNOutput, self).compute_output_shape(input_shape)


class BoschNNConv1D(Conv1D):
#     @interfaces.legacy_conv1d_support
    def __init__(self, **kwargs):
        super(BoschNNConv1D, self).__init__(**kwargs)

    def build(self, input_shape):
        super(BoschNNConv1D, self).build(input_shape)

    def call(self, inputs):
        nanMask = T.isnan(inputs)
        notNanMask =T.invert(nanMask)
        inputs_copy = T.set_subtensor(inputs[nanMask.nonzero()], 0)
        return super(BoschNNConv1D, self).call(inputs_copy) 
        
    def compute_output_shape(self, input_shape):
        return super(BoschNNConv1D, self).compute_output_shape(input_shape)

    def get_config(self):
        config = super(BoschNNConv1D, self).get_config()
        return config

# Now let's test the model with dummy data
# 
if __name__ == '__main__':
# Create empty model
    boschModel = BoschNNModel()
# Create Input Layer
    inputLayer = BoschNNInput(boschModel, 32, input_shape=(16,), name="first", activation='sigmoid')
# Create Output Layer
    outputLayer = BoschNNOutput(boschModel, 2, name='last', activation='sigmoid')

# Add layers to model
    boschModel.add(inputLayer)
#     boschModel.add(Dense(32, name='middle1', activation='sigmoid'))
#     boschModel.add(Dense(16, name='middle2', activation='sigmoid'))    
    boschModel.add(outputLayer)

# Compile model
    boschModel.compile(loss="mse", optimizer=SGD())
    
# this is what we have in the model
    boschModel.summary()

# If you want to modify initial weights.
# This is for weight matrix and bias    
    # for l in boschModel.layers:
    #     w1 = l.get_weights()
    #     print l.name, w1[0].shape, w1[1].shape
    #     w1[0] = np.ones(w1[0].shape)
    #     w1[1] = np.ones(w1[1].shape)
    #     l.set_weights(w1)
    
# Create random data b/w [-1,1)   
    np.random.seed(42)
    dummyDataset = np.random.random_sample((10000,18), ) # 16 + 2
#     dummyDataset = np.ones((10000,18)) # 16 + 2
    
    X = dummyDataset[:,:-2]
    y = dummyDataset[:,-2:]

# Add some np.NaN to the dataset    
    X[0,0] = np.NaN
    X[0,5] = np.NaN
    X[1,1] = np.NaN
    X[2,12] = np.NaN
    X[3,12] = np.NaN
    X[43,12] = np.NaN
    X[31,13] = np.NaN
    X[32,1] = np.NaN
    X[33,13] = np.NaN

# Initialize Callback for weight processing    
    boschNNCallback = BoschNNCallback()
    
    boschModel.fit(X, y, epochs=10, batch_size=1, shuffle=False, verbose=True, callbacks=[boschNNCallback])
    
# Make some predictions    
    dummyDataset = np.ones((10,16)) # 16
    dummyDataset[0,0] = np.NaN
    dummyDataset[0,5] = np.NaN
    dummyDataset[1,1] = np.NaN
    dummyDataset[2,12] = np.NaN
    dummyDataset[3,12] = np.NaN
    print boschModel.predict(dummyDataset)
    
    
#     # intermediate_tensor_function = K.function([boschModel.layers[0].input],[boschModel.layers[0].output])
#     
#     def eskiz():
#         input_tensor = boschModel.layers[0].input
#         output_var = T.isnan(input_tensor)
#         
#         output_var2 = T.set_subtensor(input_tensor[output_var.nonzero()], [4, 4, 5])
#         
#         intermediate_tensor_function = K.function([boschModel.layers[0].input],[output_var2])
#         intermediate_tensor = intermediate_tensor_function([X])[0]
#         
#         print intermediate_tensor
#     
#     # eskiz()
