import paddle.fluid as fluid


def loss(x, y, clip_value=10.0):
    """Calculate the sigmoid cross entropy with logits for input(x).

    Args:
        x: Variable with shape with shape [batch, dim]
        y: Input label

    Returns:
        loss: cross entropy
        logits: prediction
    """

    logits = fluid.layers.fc(
        input=x,
        size=1,
        bias_attr=fluid.ParamAttr(initializer=fluid.initializer.Constant(0.)))
    loss = fluid.layers.sigmoid_cross_entropy_with_logits(x=logits, label=y)
    loss = fluid.layers.reduce_mean(
        fluid.layers.clip(
            loss, min=-clip_value, max=clip_value))
    return loss, logits


def ffn(input, d_inner_hid, d_hid, name=None):
    """Position-wise Feed-Forward Network
    """

    hidden = fluid.layers.fc(input=input,
                             size=d_inner_hid,
                             num_flatten_dims=2,
                             param_attr=fluid.ParamAttr(name=name + '_fc.w_0'),
                             bias_attr=fluid.ParamAttr(
                                 name=name + '_fc.b_0',
                                 initializer=fluid.initializer.Constant(0.)),
                             act="relu")
    out = fluid.layers.fc(input=hidden,
                          size=d_hid,
                          num_flatten_dims=2,
                          param_attr=fluid.ParamAttr(name=name + '_fc.w_1'),
                          bias_attr=fluid.ParamAttr(
                              name=name + '_fc.b_1',
                              initializer=fluid.initializer.Constant(0.)))
    return out


def dot_product_attention(query,
                          key,
                          value,
                          d_key,
                          q_mask=None,
                          k_mask=None,
                          dropout_rate=None,
                          mask_cache=None):
    """Dot product layer.

     Args:
         query: a tensor with shape [batch, Q_time, Q_dimension]
         key: a tensor with shape [batch, time, K_dimension]
         value: a tensor with shape [batch, time, V_dimension]

         q_lengths: a tensor with shape [batch]
         k_lengths: a tensor with shape [batch]

     Returns:
         a tensor with shape [batch, query_time, value_dimension]

     Raises:
         AssertionError: if Q_dimension not equal to K_dimension when attention 
                        type is dot.
    """

    logits = fluid.layers.matmul(
        x=query, y=key, transpose_y=True, alpha=d_key**(-0.5))

    if (q_mask is not None) and (k_mask is not None):
        if mask_cache is not None and q_mask.name in mask_cache and k_mask.name in mask_cache[
                q_mask.name]:
            mask, another_mask = mask_cache[q_mask.name][k_mask.name]
        else:
            mask = fluid.layers.matmul(x=q_mask, y=k_mask, transpose_y=True)
            another_mask = fluid.layers.scale(
                mask,
                scale=float(2**32 - 1),
                bias=float(-1),
                bias_after_scale=False)
            if mask_cache is not None:
                if q_mask.name not in mask_cache:
                    mask_cache[q_mask.name] = dict()

                mask_cache[q_mask.name][k_mask.name] = [mask, another_mask]

        logits = mask * logits + another_mask

    attention = fluid.layers.softmax(logits)
    if dropout_rate:
        attention = fluid.layers.dropout(
            input=attention, dropout_prob=dropout_rate, is_test=False, seed=2)

    atten_out = fluid.layers.matmul(x=attention, y=value)

    return atten_out


def block(name,
          query,
          key,
          value,
          d_key,
          q_mask=None,
          k_mask=None,
          is_layer_norm=True,
          dropout_rate=None,
          mask_cache=None):
    """
    """

    att_out = dot_product_attention(
        query,
        key,
        value,
        d_key,
        q_mask,
        k_mask,
        dropout_rate,
        mask_cache=mask_cache)

    y = query + att_out
    if is_layer_norm:
        y = fluid.layers.layer_norm(
            input=y,
            begin_norm_axis=len(y.shape) - 1,
            param_attr=fluid.ParamAttr(
                initializer=fluid.initializer.Constant(1.),
                name=name + '_layer_norm.w_0'),
            bias_attr=fluid.ParamAttr(
                initializer=fluid.initializer.Constant(0.),
                name=name + '_layer_norm.b_0'))

    z = ffn(y, d_key, d_key, name)
    w = y + z
    if is_layer_norm:
        w = fluid.layers.layer_norm(
            input=w,
            begin_norm_axis=len(w.shape) - 1,
            param_attr=fluid.ParamAttr(
                initializer=fluid.initializer.Constant(1.),
                name=name + '_layer_norm.w_1'),
            bias_attr=fluid.ParamAttr(
                initializer=fluid.initializer.Constant(0.),
                name=name + '_layer_norm.b_1'))

    return w


def cnn_3d(input, out_channels_0, out_channels_1, add_relu=True):
    # same padding
    conv_0 = fluid.layers.conv3d(
        name="conv3d_0",
        input=input,
        num_filters=out_channels_0,
        filter_size=[3, 3, 3],
        padding=[1, 1, 1],
        act="elu" if add_relu else None,
        param_attr=fluid.ParamAttr(initializer=fluid.initializer.Uniform(
            low=-0.01, high=0.01)),
        bias_attr=fluid.ParamAttr(
            initializer=fluid.initializer.Constant(value=0.0)))

    # same padding 
    pooling_0 = fluid.layers.pool3d(
        input=conv_0,
        pool_type="max",
        pool_size=3,
        pool_padding=1,
        pool_stride=3)

    conv_1 = fluid.layers.conv3d(
        name="conv3d_1",
        input=pooling_0,
        num_filters=out_channels_1,
        filter_size=[3, 3, 3],
        padding=[1, 1, 1],
        act="elu" if add_relu else None,
        param_attr=fluid.ParamAttr(initializer=fluid.initializer.Uniform(
            low=-0.01, high=0.01)),
        bias_attr=fluid.ParamAttr(
            initializer=fluid.initializer.Constant(value=0.0)))

    # same padding 
    pooling_1 = fluid.layers.pool3d(
        input=conv_1,
        pool_type="max",
        pool_size=3,
        pool_padding=1,
        pool_stride=3)

    return pooling_1

[((20, 11, 20, 76), 'paddle.fluid.layers.sigmoid_cross_entropy_with_logits', 'fluid.layers.sigmoid_cross_entropy_with_logits', (), '', True, 'import paddle.fluid as fluid\n'), ((75, 13, 76, 62), 'paddle.fluid.layers.matmul', 'fluid.layers.matmul', (), '', True, 'import paddle.fluid as fluid\n'), ((97, 16, 97, 44), 'paddle.fluid.layers.softmax', 'fluid.layers.softmax', ({(97, 37, 97, 43): 'logits'}, {}), '(logits)', True, 'import paddle.fluid as fluid\n'), ((102, 16, 102, 57), 'paddle.fluid.layers.matmul', 'fluid.layers.matmul', (), '', True, 'import paddle.fluid as fluid\n'), ((173, 16, 178, 22), 'paddle.fluid.layers.pool3d', 'fluid.layers.pool3d', (), '', True, 'import paddle.fluid as fluid\n'), ((193, 16, 198, 22), 'paddle.fluid.layers.pool3d', 'fluid.layers.pool3d', (), '', True, 'import paddle.fluid as fluid\n'), ((22, 8, 23, 50), 'paddle.fluid.layers.clip', 'fluid.layers.clip', (), '', True, 'import paddle.fluid as fluid\n'), ((99, 20, 100, 78), 'paddle.fluid.layers.dropout', 'fluid.layers.dropout', (), '', True, 'import paddle.fluid as fluid\n'), ((34, 40, 34, 78), 'paddle.fluid.ParamAttr', 'fluid.ParamAttr', (), '', True, 'import paddle.fluid as fluid\n'), ((42, 37, 42, 75), 'paddle.fluid.ParamAttr', 'fluid.ParamAttr', (), '', True, 'import paddle.fluid as fluid\n'), ((83, 19, 83, 76), 'paddle.fluid.layers.matmul', 'fluid.layers.matmul', (), '', True, 'import paddle.fluid as fluid\n'), ((19, 46, 19, 76), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(19, 73, 19, 75): '0.0'}, {}), '(0.0)', True, 'import paddle.fluid as fluid\n'), ((37, 45, 37, 75), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(37, 72, 37, 74): '0.0'}, {}), '(0.0)', True, 'import paddle.fluid as fluid\n'), ((45, 42, 45, 72), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(45, 69, 45, 71): '0.0'}, {}), '(0.0)', True, 'import paddle.fluid as fluid\n'), ((167, 47, 168, 33), 'paddle.fluid.initializer.Uniform', 'fluid.initializer.Uniform', (), '', True, 'import paddle.fluid as fluid\n'), ((170, 24, 170, 61), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', (), '', True, 'import paddle.fluid as fluid\n'), ((187, 47, 188, 33), 'paddle.fluid.initializer.Uniform', 'fluid.initializer.Uniform', (), '', True, 'import paddle.fluid as fluid\n'), ((190, 24, 190, 61), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', (), '', True, 'import paddle.fluid as fluid\n'), ((136, 28, 136, 58), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(136, 55, 136, 57): '1.0'}, {}), '(1.0)', True, 'import paddle.fluid as fluid\n'), ((139, 28, 139, 58), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(139, 55, 139, 57): '0.0'}, {}), '(0.0)', True, 'import paddle.fluid as fluid\n'), ((149, 28, 149, 58), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(149, 55, 149, 57): '1.0'}, {}), '(1.0)', True, 'import paddle.fluid as fluid\n'), ((152, 28, 152, 58), 'paddle.fluid.initializer.Constant', 'fluid.initializer.Constant', ({(152, 55, 152, 57): '0.0'}, {}), '(0.0)', True, 'import paddle.fluid as fluid\n')]