import lasagne as nn
import theano.tensor as T
import numpy as np
from lasagne import nonlinearities
from lasagne.layers.dnn import Conv2DDNNLayer


def lb_softplus(lb=1):
    return lambda x: nn.nonlinearities.softplus(x) + lb


class MultLayer(nn.layers.MergeLayer):
    """
    takes elementwise product between 2 layers
    """

    def __init__(self, input1, input2, log=False, **kwargs):
        super(MultLayer, self).__init__([input1, input2], **kwargs)

    def get_output_shape_for(self, input_shapes):
        return input_shapes[0]

    def get_output_for(self, inputs, **kwargs):
        return inputs[0] * inputs[1]


class ConstantLayer(nn.layers.Layer):
    """
    Makes a layer of constant value the same shape as the given input layer
    """

    def __init__(self, shape_layer, constant=1, **kwargs):
        super(ConstantLayer, self).__init__(shape_layer, **kwargs)
        self.constant = constant

    def get_output_shape_for(self, input_shape):
        return input_shape

    def get_output_for(self, input, **kwargs):
        return T.ones_like(input) * self.constant


class RepeatLayer(nn.layers.Layer):
    def __init__(self, incoming, repeats, axis=0, **kwargs):
        super(RepeatLayer, self).__init__(incoming, **kwargs)
        self.repeats = repeats
        self.axis = axis

    def get_output_shape_for(self, input_shape):
        output_shape = list(input_shape)
        output_shape.insert(self.axis, self.repeats)
        return tuple(output_shape)

    def get_output_for(self, input, **kwargs):
        shape_ones = [1] * input.ndim
        shape_ones.insert(self.axis, self.repeats)
        ones = T.ones(tuple(shape_ones), dtype=input.dtype)

        pattern = range(input.ndim)
        pattern.insert(self.axis, "x")
        # print shape_ones, pattern
        return ones * input.dimshuffle(*pattern)


class AttentionLayer(nn.layers.Layer):
    def __init__(self, incoming, u=nn.init.GlorotUniform(), **kwargs):
        super(AttentionLayer, self).__init__(incoming, **kwargs)
        num_inputs = self.input_shape[-1]
        self.u = self.add_param(u, (num_inputs, 1), name='u')

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

    def get_output_for(self, input, **kwargs):
        a = T.nnet.softmax(T.dot(input, self.u)[:, :, 0])
        return T.sum(a[:, :, np.newaxis] * input, axis=1)


class MaskedMeanPoolLayer(nn.layers.MergeLayer):
    """
    pools globally across all trailing dimensions beyond the given axis.
    give it a mask
    """

    def __init__(self, incoming, mask, axis, **kwargs):
        super(MaskedMeanPoolLayer, self).__init__([incoming, mask], **kwargs)
        self.axis = axis

    def get_output_shape_for(self, input_shapes):
        return input_shapes[0][:self.axis] + (1,)

    def get_output_for(self, inputs, **kwargs):
        input = inputs[0]
        mask = inputs[1]
        masked_input = input * mask.dimshuffle(0, 1, 'x')
        return T.sum(masked_input.flatten(self.axis + 1), axis=self.axis, keepdims=True) / T.sum(mask, axis=-1,
                                                                                                 keepdims=True)


class MaskedSTDPoolLayer(nn.layers.MergeLayer):
    """
    pools globally across all trailing dimensions beyond the given axis.
    give it a mask
    """

    def __init__(self, incoming, mask, axis, **kwargs):
        super(MaskedSTDPoolLayer, self).__init__([incoming, mask], **kwargs)
        self.axis = axis

    def get_output_shape_for(self, input_shapes):
        return input_shapes[0][:self.axis] + (1,)

    def get_output_for(self, inputs, **kwargs):
        input = inputs[0]
        mask = inputs[1]
        masked_input = input * mask.dimshuffle(0, 1, 'x')
        mu_x = T.sum(masked_input.flatten(self.axis + 1), axis=self.axis, keepdims=True) / T.sum(mask, axis=-1,
                                                                                                 keepdims=True)
        mu_x2 = T.sum(masked_input.flatten(self.axis + 1) ** 2, axis=self.axis, keepdims=True) / T.sum(mask, axis=-1,
                                                                                                       keepdims=True)
        return T.sqrt(mu_x2 - mu_x ** 2)


class NonlinearityLayer(nn.layers.Layer):
    def __init__(self, incoming, nonlinearity=nonlinearities.rectify,
                 **kwargs):
        super(NonlinearityLayer, self).__init__(incoming, **kwargs)
        self.nonlinearity = (nonlinearities.identity if nonlinearity is None
                             else nonlinearity)

    def get_output_for(self, input, **kwargs):
        return self.nonlinearity(input)


class CumSumLayer(nn.layers.Layer):
    def __init__(self, incoming, axis=1, **kwargs):
        super(CumSumLayer, self).__init__(incoming, **kwargs)
        self.axis = axis

    def get_output_shape_for(self, input_shape):
        return input_shape

    def get_output_for(self, input, **kwargs):
        result = T.extra_ops.cumsum(input, axis=self.axis)
        return result


class NormalisationLayer(nn.layers.Layer):
    def __init__(self, incoming, norm_sum=1.0, allow_negative=False, **kwargs):
        super(NormalisationLayer, self).__init__(incoming, **kwargs)
        self.norm_sum = norm_sum
        self.allow_negative = allow_negative

    def get_output_for(self, input, **kwargs):
        # take the minimal working slice size, and use that one.
        if self.allow_negative:
            inp_low_zero = input - T.min(input, axis=1).dimshuffle(0, 'x')
        else:
            inp_low_zero = input
        return inp_low_zero / T.sum(inp_low_zero, axis=1).dimshuffle(0, 'x') * self.norm_sum


class HighwayLayer(nn.layers.MergeLayer):
    def __init__(self, gate, input1, input2, **kwargs):
        incomings = [gate, input1, input2]
        super(HighwayLayer, self).__init__(incomings, **kwargs)
        assert gate.output_shape == input1.output_shape == input2.output_shape

    def get_output_shape_for(self, input_shapes):
        return input_shapes[0]

    def get_output_for(self, inputs, **kwargs):
        return inputs[0] * inputs[1] + (1 - inputs[0]) * inputs[2]


def highway_conv3(incoming, nonlinearity=nn.nonlinearities.rectify, **kwargs):
    wh = nn.init.Orthogonal('relu')
    bh = nn.init.Constant(0.0)
    wt = nn.init.Orthogonal('relu')
    bt = nn.init.Constant(-2.)
    num_filters = incoming.output_shape[1]

    # H
    l_h = Conv2DDNNLayer(incoming, num_filters=num_filters,
                         filter_size=(3, 3), stride=(1, 1),
                         pad='same', W=wh, b=bh,
                         nonlinearity=nonlinearity)
    # T
    l_t = Conv2DDNNLayer(incoming, num_filters=num_filters,
                         filter_size=(3, 3), stride=(1, 1),
                         pad='same', W=wt, b=bt,
                         nonlinearity=T.nnet.sigmoid)

    return HighwayLayer(gate=l_t, input1=l_h, input2=incoming)


class Upscale3DLayer(nn.layers.Layer):
    """
    3D upscaling layer
    Performs 3D upscaling over the three trailing axes of a 5D input tensor.
    Parameters
    ----------
    incoming : a :class:`Layer` instance or tuple
        The layer feeding into this layer, or the expected input shape.
    scale_factor : integer or iterable
        The scale factor in each dimension. If an integer, it is promoted to
        a cubic scale factor region. If an iterable, it should have three
        elements.
    mode : {'repeat', 'dilate'}
        Upscaling mode: repeat element values or upscale leaving zeroes between
        upscaled elements. Default is 'repeat'.
    **kwargs
        Any additional keyword arguments are passed to the :class:`Layer`
        superclass.
    """

    def __init__(self, incoming, scale_factor, mode='repeat', **kwargs):
        super(Upscale3DLayer, self).__init__(incoming, **kwargs)

        self.scale_factor = nn.utils.as_tuple(scale_factor, 3)

        if self.scale_factor[0] < 1 or self.scale_factor[1] < 1 or \
                        self.scale_factor[2] < 1:
            raise ValueError('Scale factor must be >= 1, not {0}'.format(
                self.scale_factor))

        if mode not in {'repeat', 'dilate'}:
            msg = "Mode must be either 'repeat' or 'dilate', not {0}"
            raise ValueError(msg.format(mode))
        self.mode = mode

    def get_output_shape_for(self, input_shape):
        output_shape = list(input_shape)  # copy / convert to mutable list
        if output_shape[2] is not None:
            output_shape[2] *= self.scale_factor[0]
        if output_shape[3] is not None:
            output_shape[3] *= self.scale_factor[1]
        if output_shape[4] is not None:
            output_shape[4] *= self.scale_factor[2]
        return tuple(output_shape)

    def get_output_for(self, input, **kwargs):
        a, b, c = self.scale_factor
        upscaled = input
        if self.mode == 'repeat':
            if c > 1:
                upscaled = T.extra_ops.repeat(upscaled, c, 4)
            if b > 1:
                upscaled = T.extra_ops.repeat(upscaled, b, 3)
            if a > 1:
                upscaled = T.extra_ops.repeat(upscaled, a, 2)
        elif self.mode == 'dilate':
            if c > 1 or b > 1 or a > 1:
                output_shape = self.get_output_shape_for(input.shape)
                upscaled = T.zeros(shape=output_shape, dtype=input.dtype)
                upscaled = T.set_subtensor(
                    upscaled[:, :, ::a, ::b, ::c], input)
        return upscaled


class CastingLayer(nn.layers.Layer):
    def __init__(self, incoming, dtype, **kwargs):
        super(CastingLayer, self).__init__(incoming, **kwargs)
        self.dtype = dtype

    def get_output_for(self, input, **kwargs):
        return T.cast(input, self.dtype)


def heaviside(x, size):
    return T.arange(0, size).dimshuffle('x', 0) - T.repeat(x, size, axis=1) >= 0.


class NormalCDFLayer(nn.layers.MergeLayer):
    def __init__(self, mu, sigma, max_support, **kwargs):
        super(NormalCDFLayer, self).__init__([mu, sigma], **kwargs)
        self.max_support = max_support

    def get_output_shape_for(self, input_shapes):
        return input_shapes[0][0], self.max_support

    def get_output_for(self, input, **kwargs):
        mu = input[0]
        sigma = input[1]

        x_range = T.arange(0, self.max_support).dimshuffle('x', 0)
        mu = T.repeat(mu, self.max_support, axis=1)
        sigma = T.repeat(sigma, self.max_support, axis=1)
        x = (x_range - mu) / (sigma * T.sqrt(2.) + 1e-16)
        cdf = (T.erf(x) + 1.) / 2.
        return cdf



class AggAllBenignExp(nn.layers.Layer):
    """
    Aggregates the chances
    """

    def __init__(self, incoming, **kwargs):
        super(AggAllBenignExp, self).__init__(incoming, **kwargs)

    def get_output_shape_for(self, input_shape):
        assert(len(input_shape)==3)
        assert(input_shape[2]==1)
        return (input_shape[0], 1)

    def get_output_for(self, input, **kwargs):
        rectified = nonlinearities.softplus(input)
        sum_rect = T.sum(rectified, axis=(1,2))
        output = 1 - T.exp(-sum_rect)
        return output

class AggAllBenignProd(nn.layers.Layer):
    """
    takes elementwise product between 2 layers
    """

    def __init__(self, incoming, apply_nl = True, **kwargs):
        super(AggAllBenignProd, self).__init__(incoming, **kwargs)
        self.apply_nl = apply_nl

    def get_output_shape_for(self, input_shape):
        assert(len(input_shape)==3)
        assert(input_shape[2]==1)
        return (input_shape[0], 1)

    def get_output_for(self, input, **kwargs):
        if apply_nl:
            ps = nonlinearities.sigmoid(input)
        prod = T.prod(ps, axis=(1,2))
        output = 1 - prod
        return output

class AggSoPP(nn.layers.Layer):
    """
    Aggregates via Sum of powers
    """
    def __init__(self, incoming, exp=nn.init.Constant(2.),  **kwargs):
        super(AggSoPP, self).__init__(incoming, **kwargs)
        self.exp = self.add_param(exp, (1,), name='exp', regularizable=False)

    def get_output_shape_for(self, input_shape):
        assert(len(input_shape)==3)
        assert(input_shape[2]==1)
        return (input_shape[0], 1)

    def get_output_for(self, input, **kwargs):
        ps = nonlinearities.sigmoid(input)
        powd = ps ** self.exp
        tmean = T.mean(powd, axis=(1,2))
        return tmean

class Unbroadcast(nn.layers.Layer):
    """
    takes elementwise product between 2 layers
    """
    def __init__(self, incoming,  **kwargs):
        super(Unbroadcast, self).__init__(incoming, **kwargs)

    def get_output_shape_for(self, input_shape):
        return input_shape

    def get_output_for(self, input, **kwargs):
        all_dims = range(len(T.shape(input)))
        print all_dims
        return T.Unbroadcast(input, *all_dims)

class LogMeanExp(nn.layers.Layer):
    """
    ln(mean(exp( r * x ))) /  r
    """

    def __init__(self, incoming, r=1, axis=-1, **kwargs):
        super(LogMeanExp, self).__init__(incoming, **kwargs)
        self.r = np.float32(r)
        self.axis = axis

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

    def get_output_for(self, input, **kwargs):
        return T.log(T.mean(T.exp(self.r * input), axis=self.axis) + 1e-7) / self.r


class AggMILLoss(nn.layers.Layer):
    """
    ln(mean(exp( r * x ))) /  r
    """

    def __init__(self, incoming, r=1, **kwargs):
        super(AggMILLoss, self).__init__(incoming, **kwargs)
        self.r = np.float32(r)

    def get_output_shape_for(self, input_shape):
        assert(len(input_shape)==3)
        assert(input_shape[2]==1)
        return (input_shape[0], 2)

    def get_output_for(self, input, **kwargs):

        ps = nonlinearities.sigmoid(input)
        sum_p_r_benign = T.sum(ps,axis=1)
        sum_log = T.sum(T.log(1-ps+1.e-12),axis=1)
        return T.concatenate([sum_log, sum_p_r_benign])


class Real3DFFTLayer(nn.layers.Layer):
    """
    Real3DFFTLayer
    """

    def __init__(self, incoming, r=1, **kwargs):
        super(Real3DFFTLayer, self).__init__(incoming, **kwargs)

    def get_output_shape_for(self, input_shape):
        assert(len(input_shape)==4)
        output_shape = (input_shape[0],)
        for i in range(1,len(input_shape)):
            output_shape = output_shape + (input_shape[i]//2,)        
        return output_shape

    def get_output_for(self, input, **kwargs):
        cfft = T.fft.rfft(input)
        magnitude = (cfft[:,:,:,0]**2 + cfft[:,:,:,1]**2)**0.5
        return magnitude