# %% Model definition + helper functions
# %%

# Import modules and libraries
from keras.models import Model
from keras.layers import Input, Activation, UpSampling2D, Convolution2D, MaxPooling2D, BatchNormalization
from keras.callbacks import Callback
from keras import backend as K
from keras import optimizers
from keras import losses
import tensorflow as tf
import numpy as np


# define a function that projects an image to the range [0,1]
def project_01(im):
    im = np.squeeze(im)
    min_val = im.min()
    max_val = im.max()
    return (im - min_val) / (max_val - min_val)


# normalize image given mean and std
def normalize_im(im, dmean, dstd):
    im = np.squeeze(im)
    im_norm = np.zeros(im.shape, dtype=np.float32)
    im_norm = (im - dmean) / dstd
    return im_norm


# Define the loss history recorder
class LossHistory(Callback):
    def on_train_begin(self, logs={}):
        self.losses = []

    def on_batch_end(self, batch, logs={}):
        self.losses.append(logs.get('loss'))


#  Define a matlab like gaussian 2D filter
def matlab_style_gauss2D(shape=(7, 7), sigma=1):
    """
    2D gaussian filter - should give the same result as:
    MATLAB's fspecial('gaussian',[shape],[sigma])
    """
    m, n = [(ss - 1.) / 2. for ss in shape]
    y, x = np.ogrid[-m:m + 1, -n:n + 1]
    h = np.exp(-(x * x + y * y) / (2. * sigma * sigma))
    h.astype(dtype=K.floatx())
    h[h < np.finfo(h.dtype).eps * h.max()] = 0
    sumh = h.sum()
    if sumh != 0:
        h /= sumh
    h = h * 2.0
    h = h.astype('float32')
    return h


# Expand the filter dimensions
psf_heatmap = matlab_style_gauss2D(shape=(7, 7), sigma=1)
gfilter = tf.reshape(psf_heatmap, [7, 7, 1, 1])


# Combined MSE + L1 loss
def L1L2loss(input_shape):
    def bump_mse(heatmap_true, spikes_pred):
        # generate the heatmap corresponding to the predicted spikes
        heatmap_pred = K.conv2d(spikes_pred, gfilter, strides=(1, 1), padding='same')

        # heatmaps MSE
        loss_heatmaps = losses.mean_squared_error(heatmap_true, heatmap_pred)

        # l1 on the predicted spikes
        loss_spikes = losses.mean_absolute_error(spikes_pred, tf.zeros(input_shape))
        return loss_heatmaps + loss_spikes

    return bump_mse


# Define the concatenated conv2, batch normalization, and relu block
def conv_bn_relu(nb_filter, rk, ck, name):
    def f(input):
        conv = Convolution2D(nb_filter, kernel_size=(rk, ck), strides=(1, 1), \
                             padding="same", use_bias=False, \
                             kernel_initializer="Orthogonal", name='conv-' + name)(input)
        conv_norm = BatchNormalization(name='BN-' + name)(conv)
        conv_norm_relu = Activation(activation="relu", name='Relu-' + name)(conv_norm)
        return conv_norm_relu

    return f


# Define the model architechture
def CNN(input, names):
    Features1 = conv_bn_relu(32, 3, 3, names + 'F1')(input)
    pool1 = MaxPooling2D(pool_size=(2, 2), name=names + 'Pool1')(Features1)
    Features2 = conv_bn_relu(64, 3, 3, names + 'F2')(pool1)
    pool2 = MaxPooling2D(pool_size=(2, 2), name=names + 'Pool2')(Features2)
    Features3 = conv_bn_relu(128, 3, 3, names + 'F3')(pool2)
    pool3 = MaxPooling2D(pool_size=(2, 2), name=names + 'Pool3')(Features3)
    Features4 = conv_bn_relu(512, 3, 3, names + 'F4')(pool3)
    up5 = UpSampling2D(size=(2, 2), name=names + 'Upsample1')(Features4)
    Features5 = conv_bn_relu(128, 3, 3, names + 'F5')(up5)
    up6 = UpSampling2D(size=(2, 2), name=names + 'Upsample2')(Features5)
    Features6 = conv_bn_relu(64, 3, 3, names + 'F6')(up6)
    up7 = UpSampling2D(size=(2, 2), name=names + 'Upsample3')(Features6)
    Features7 = conv_bn_relu(32, 3, 3, names + 'F7')(up7)
    return Features7


# Define the Model building for an arbitrary input size
def buildModel(input_dim):
    input_ = Input(shape=(input_dim),name="CNN_input")
    act_ = CNN(input_, 'CNN')
    density_pred = Convolution2D(1, kernel_size=(1, 1), strides=(1, 1), padding="same", \
                                 activation="linear", use_bias=False, \
                                 kernel_initializer="Orthogonal", name='Prediction')(act_)
    model = Model(inputs=input_, outputs=density_pred)
    opt = optimizers.Adam(lr=0.001)
    model.compile(optimizer=opt, loss=L1L2loss(input_dim))
    return model
