import tensorflow as tf
from openrec.modules.extractions import MultiLayerFC


def MLPSoftmax(
    user,
    item,
    seq_len,
    max_seq_len,
    dims,
    subgraph,
    item_bias=None,
    extra=None,
    l2_reg=None,
    labels=None,
    dropout=None,
    train=None,
    scope=None,
):

    with tf.variable_scope(scope, reuse=tf.AUTO_REUSE):

        # average item vectors user interacted with
        seq_mask = tf.sequence_mask(seq_len, max_seq_len, dtype=tf.float32)
        item = tf.reduce_mean(item * tf.expand_dims(seq_mask, axis=2), axis=1)

        if user is not None:
            in_tensor = tf.concat([user, item], axis=1)
        else:
            in_tensor = tf.concat([item], axis=1)

        if extra is not None:
            in_tensor = tf.concat([in_tensor, extra], axis=1)

        if train:
            logits = MultiLayerFC(
                in_tensor=in_tensor,
                dims=dims,
                subgraph=subgraph,
                bias_in=True,
                bias_mid=True,
                bias_out=False,
                dropout_mid=dropout,
                l2_reg=l2_reg,
                scope="mlp_reg",
            )
        else:
            logits = MultiLayerFC(
                in_tensor=in_tensor,
                dims=dims,
                subgraph=subgraph,
                bias_in=True,
                bias_mid=True,
                bias_out=False,
                l2_reg=l2_reg,
                scope="mlp_reg",
            )

        if item_bias is not None:
            logits += item_bias

        if train:
            loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
                labels=labels, logits=logits
            )
            subgraph.register_global_loss(tf.reduce_mean(loss))
        else:
            subgraph.register_global_output(logits)
