import lasagne.layers as layers
import numpy as np
import theano.tensor as T
from lasagne import nonlinearities
from lasagne.objectives import aggregate, categorical_crossentropy, squared_error
from lasagne.regularization import l1, regularize_layer_params
from nolearn.lasagne import NeuralNet

########################################################
## Load the "Dataset"
########################################################
N = 100
data = np.zeros((N, 3, 128, 128), np.float32)  # Fake Image Dataset.
cls = np.random.randint(0, 10, N).astype(np.float32)  # Classes for images
reg = np.random.random(N).astype(np.float32)  # Some regression value
targets = np.stack([cls, reg], axis=1)

########################################################
## Define The Network
########################################################
prop = dict(nonlinearity=nonlinearities.leaky_rectify)
l_in = layers.InputLayer((None, 3, 128, 128), name="input")

# Shared Convolutions
l_conv1 = layers.Conv2DLayer(l_in, num_filters=32, filter_size=5, name="shared_conv1", **prop)
l_pool1 = layers.Pool2DLayer(l_conv1, pool_size=2, name="shared_pool1")
l_conv2 = layers.Conv2DLayer(l_pool1, num_filters=64, filter_size=3, name="shared_conv2", **prop)
l_pool2 = layers.Pool2DLayer(l_conv2, pool_size=2, name="shared_pool2")
l_conv3 = layers.Conv2DLayer(l_pool2, num_filters=128, filter_size=3, name="shared_conv3", **prop)
l_pool3 = layers.Pool2DLayer(l_conv3, pool_size=2, name="shared_pool3")

# Task 1 - 10 Category Classifier
l_tsk1 = layers.DenseLayer(l_pool3, name="task1_d1", num_units=45, **prop)
l_tsk1 = layers.DenseLayer(l_tsk1, name="task1_d2", num_units=64, **prop)
l_tsk1 = layers.DenseLayer(l_tsk1, name="task1_out", num_units=10, nonlinearity=nonlinearities.softmax)

# Task 2 -  1 Regression
l_tsk2 = layers.DenseLayer(l_pool3, name="task2_d1", num_units=54, **prop)
l_tsk2 = layers.DenseLayer(l_tsk2, name="task2_d2", num_units=64, **prop)
l_tsk2 = layers.DenseLayer(l_tsk2, name="task2_out", num_units=1, nonlinearity=nonlinearities.sigmoid)

out_layers = [l_tsk1, l_tsk2]


########################################################
## Write a custom Objective
########################################################
def cls_regression(layers_, target, **kwargs):
    output_kw = kwargs.pop('get_output_kw', {})
    det = kwargs.pop('deterministic', False)

    l1_penalty = kwargs.pop("l1", 0)
    cls_lmb = kwargs.pop("cls_lambda", 1)
    reg_lmb = kwargs.pop("reg_lambda", 1)

    cls_layer = layers_['task1_out']
    reg_layer = layers_['task2_out']

    # Get the outputs
    out_cls, out_reg = layers.get_output([cls_layer, reg_layer], deterministic=det, **output_kw)

    # Get the targets
    gt_cls = T.cast(target[:, 0], 'int32')
    gt_reg = target[:, 1].reshape((-1, 1))

    # Calculate the multi task loss
    cls_loss = cls_lmb * aggregate(categorical_crossentropy(out_cls, gt_cls))
    reg_loss = reg_lmb * aggregate(squared_error(out_reg, gt_reg))
    loss = cls_loss + reg_loss

    if l1_penalty:
        loss += l1_penalty * regularize_layer_params(layers_.values(), l1)

    return loss


########################################################
## Instantiate the Network The Network
########################################################
network = NeuralNet(layers=out_layers,

                    regression=True,  # <=== Its probably easier to leave as regression, and do any argmax manually
                    custom_scores=[],  # <=== Custom scores will only be passed the output of out_layers[-1]
                    y_tensor_type=None,  # <===Output type of out_layers[-1]
                    max_epochs=200,
                    objective=cls_regression,
                    update_learning_rate=0.1,
                    update_momentum=0.9,
                    verbose=True,
                    )

# fit will correctly treat as a multi output network.
# Everything else will treat as a single output network with output = out_layers[-1]
network.fit(data, targets,epochs=1)

# This will only score out_layers[-1], targets must only include values for out_layers[-1]
sc = network.score(data[10:15], targets[10:15,1])
print(sc)

# This output here will only include outputs from out_layers[-1]
p = network.predict(data[10:15])
print(p.shape)