<reponame>caglar/Attentive_reader
from commons import EPS
import theano.tensor as TT
from utils import as_floatX
from theano.tensor.extra_ops import cumsum
import theano


def kl_simple(Y, Y_hat, cost_mask=None):
    """
    Warning: This function expects a sigmoid nonlinearity in the
    output layer. Returns a batch (vector) of mean across units of
    KL divergence for each example,
    KL(P || Q) where P is defined by Y and Q is defined by Y_hat:
    p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)
    For binary p, some terms drop out:
    - p log q - (1-p) log (1-q)
    - p log sigmoid(z) - (1-p) log sigmoid(-z)
    p softplus(-z) + (1-p) softplus(z)
    Parameters
    ----------
    Y : Variable
        targets for the sigmoid outputs. Currently Y must be purely binary.
        If it's not, you'll still get the right gradient, but the
        value in the monitoring channel will be wrong.
    Y_hat : Variable
        predictions made by the sigmoid layer. Y_hat must be generated by
        fprop, i.e., it must be a symbolic sigmoid.
    -------
    ave : Variable
        average kl divergence between Y and Y_hat.
    """

    assert hasattr(Y_hat, 'owner')

    owner = Y_hat.owner
    assert owner is not None
    op = owner.op

    if not hasattr(op, 'scalar_op'):
        raise ValueError("Expected Y_hat to be generated by an Elemwise "
                         "op, got "+str(op)+" of type "+str(type(op)))

    assert isinstance(op.scalar_op, TT.nnet.sigm.ScalarSigmoid)

    z, = owner.inputs

    z = z.reshape(Y.shape)
    term_1 = Y * TT.nnet.softplus(-z)
    term_2 = (1 - Y) * TT.nnet.softplus(z)

    total = term_1 + term_2
    if cost_mask is not None:
        total = cost_mask * total

    ave = total.sum() / (total.shape[1] * (total.shape[2] - 2))
    #ave = ave.mean()
    return ave


def kl(Y, Y_hat, cost_mask=None,
       batch_vec=True,
       cost_matrix=False,
       sum_tru_time=False,
       normalize_by_outsize=True):

    """
    Warning: This function expects a sigmoid nonlinearity in the
    output layer. Returns a batch (vector) of mean across units of
    KL divergence for each example,
    KL(P || Q) where P is defined by Y and Q is defined by Y_hat:
    p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)
    For binary p, some terms drop out:
    - p log q - (1-p) log (1-q)
    - p log sigmoid(z) - (1-p) log sigmoid(-z)
    p softplus(-z) + (1-p) softplus(z)
    Parameters
    ----------
    Y : Variable
        targets for the sigmoid outputs. Currently Y must be purely binary.
        If it's not, you'll still get the right gradient, but the
        value in the monitoring channel will be wrong.
    Y_hat : Variable
        predictions made by the sigmoid layer. Y_hat must be generated by
        fprop, i.e., it must be a symbolic sigmoid.
    -------
    ave : Variable
        average kl divergence between Y and Y_hat.
    """

    assert hasattr(Y_hat, 'owner')

    owner = Y_hat.owner
    assert owner is not None
    op = owner.op

    if not hasattr(op, 'scalar_op'):
        raise ValueError("Expected Y_hat to be generated by an Elemwise "
                         "op, got "+str(op)+" of type "+str(type(op)))

    assert isinstance(op.scalar_op, TT.nnet.sigm.ScalarSigmoid)

    z, = owner.inputs
    z = z.reshape(Y.shape)
    term_1 = Y * TT.nnet.softplus(-z)
    term_2 = (1 - Y) * TT.nnet.softplus(z)

    total = term_1 + term_2

    if cost_mask is not None:
        if cost_mask.ndim != total.ndim:
            cost_mask = cost_mask.dimshuffle(0, 1, 'x')

        total = cost_mask * total
    if not sum_tru_time:
        if normalize_by_outsize:
            if cost_matrix:
                ave = total.sum(-1) / TT.cast((total.shape[2] - 2), "float32")
            else:
                if batch_vec:
                    ave = total.sum(0).sum(1) / TT.cast((total.shape[2] - 2), "float32")
                else:
                    ave = total.sum() / (total.shape[1] * (total.shape[2] - 2))
        else:
            if cost_matrix:
                ave = total.sum(-1)
            else:
                if batch_vec:
                    ave = total.sum(0).sum(1)
                else:
                    ave = total.sum() / TT.cast(total.shape[1], "float32")
    else:
        assert not cost_matrix
        if normalize_by_outsize:
            if batch_vec:
                ave = cumsum(total.sum(-1) / TT.cast((total.shape[2] - 2), "float32"),
                        axis=0)[::-1]
            else:
                ave = cumsum(total.sum((1, 2)) / (total.shape[1] * (total.shape[2] - 2)),
                        axis=0)[::-1]
        else:
            if batch_vec:
                ave = cumsum(total.sum(-1), axis=0)[::-1]
            else:
                ave = cumsum(total.sum((1, 2)) / TT.cast(total.shape[1], "float32"), axis=0)[::-1]

    return ave


def _grab_probs(class_probs, target, use_fast_ver=False):
    if class_probs.ndim == 3:
        class_probs = class_probs.reshape((-1, class_probs.shape[-1]))

    shape0 = class_probs.shape[0]
    shape1 = class_probs.shape[1]

    p = None
    if target.ndim == 2 and use_fast_ver:
        target = target.flatten()
        cp = class_probs.reshape((target.shape[0], -1))
        p = TT.diag(cp.T[target])
    else:
        if target.ndim > 1:
            target = target.flatten()
        assert target.ndim == 1, 'make sure target is a vector of ints'
        assert 'int' in target.dtype
        pos = TT.arange(shape0)*shape1
        new_targ = target + pos
        p = class_probs.reshape((shape0*shape1, 1))[new_targ].reshape((shape0,))
    return p


def nll_simple(Y, Y_hat,
               cost_mask=None,
               cost_ent_mask=None,
               cost_ent_desc_mask=None):

    probs = Y_hat
    pred = TT.argmax(probs, axis=1).reshape(Y.shape)
    errors = TT.neq(pred, Y)
    ent_errors = None
    if cost_ent_mask is not None:
        pred_ent = TT.argmax(probs * cost_ent_mask.dimshuffle('x', 0),
                             axis=1).reshape(Y.shape)
        ent_errors = TT.neq(pred_ent, Y).mean()

    ent_desc_errors = None
    if cost_ent_desc_mask is not None:
        pred_desc_ent = TT.argmax(probs * cost_ent_desc_mask,
                             axis=1).reshape(Y.shape)
        ent_desc_errors = TT.neq(pred_desc_ent, Y).mean()

    LL = TT.log(_grab_probs(probs, Y) + 1e-8).reshape(Y.shape)

    if cost_mask is not None:
        total = cost_mask * LL
        errors = cost_mask * errors
        ncosts = TT.sum(cost_mask)
        mean_errors = TT.sum(errors) / (ncosts)
        ave = -TT.sum(total) / Y.shape[1]
    else:
        mean_errors = TT.mean(errors)
        ave = -TT.sum(LL) / Y.shape[0]
    return ave, mean_errors, ent_errors, ent_desc_errors


def nll(Y, Y_hat, cost_mask=None, batch_vec=True,
        cost_matrix=False, use_fast_ver=False):
    probs = Y_hat
    pred = TT.argmax(probs, axis=-1).reshape(Y.shape)
    errors = TT.neq(pred, Y)
    LL = TT.log(_grab_probs(probs, Y, use_fast_ver=use_fast_ver) + 1e-8).reshape(Y.shape)

    if cost_mask is not None:
        total = cost_mask * LL
        errors = cost_mask * errors
        ncosts = TT.sum(cost_mask)
        mean_errors = TT.sum(errors) / (ncosts)
        if cost_matrix:
            ave = -total
        else:
            if batch_vec:
                ave = -total.sum(0)
            else:
                ave = -TT.sum(total) / Y.shape[1]
    else:
        mean_errors = TT.mean(errors)
        if cost_matrix:
            ave = -LL
        else:
            if batch_vec:
               ave = -LL.sum(0)
            else:
               ave = -TT.sum(LL) / Y.shape[1]

    return ave, mean_errors


def multiclass_hinge_loss(predictions, targets, delta=1):
    """Computes the multi-class hinge loss between predictions and targets.
    .. math:: L_i = \\max_{j \\not = p_i} (0, t_j - t_{p_i} + \\delta)
    Parameters
    ----------
    predictions : Theano 2D tensor
        Predictions in (0, 1), such as softmax output of a neural network,
        with data points in rows and class probabilities in columns.
    targets : Theano 2D tensor or 1D tensor
        Either a vector of int giving the correct class index per data point
        or a 2D tensor of one-hot encoding of the correct class in the same
        layout as predictions (non-binary targets in [0, 1] do not work!)
    delta : scalar, default 1
        The hinge loss margin
    Returns
    -------
    Theano 1D tensor
        An expression for the item-wise multi-class hinge loss
    Notes
    -----
    This is an alternative to the categorical cross-entropy loss for
    multi-class classification problems
    """
    num_cls = predictions.shape[1]
    if targets.ndim == predictions.ndim - 1:
        targets = theano.tensor.extra_ops.to_one_hot(targets, num_cls)
    elif targets.ndim != predictions.ndim:
        raise TypeError('rank mismatch between targets and predictions')
    corrects = predictions[targets.nonzero()]
    rest = theano.tensor.reshape(predictions[(1-targets).nonzero()],
                                 (-1, num_cls-1))
    rest = theano.tensor.max(rest, axis=1)
    return theano.tensor.nnet.relu(rest - corrects + delta).mean(0).sum()

from commons import EPS
import theano.tensor as TT
from utils import as_floatX
from theano.tensor.extra_ops import cumsum
import theano


def kl_simple(Y, Y_hat, cost_mask=None):
    """
    Warning: This function expects a sigmoid nonlinearity in the
    output layer. Returns a batch (vector) of mean across units of
    KL divergence for each example,
    KL(P || Q) where P is defined by Y and Q is defined by Y_hat:
    p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)
    For binary p, some terms drop out:
    - p log q - (1-p) log (1-q)
    - p log sigmoid(z) - (1-p) log sigmoid(-z)
    p softplus(-z) + (1-p) softplus(z)
    Parameters
    ----------
    Y : Variable
        targets for the sigmoid outputs. Currently Y must be purely binary.
        If it's not, you'll still get the right gradient, but the
        value in the monitoring channel will be wrong.
    Y_hat : Variable
        predictions made by the sigmoid layer. Y_hat must be generated by
        fprop, i.e., it must be a symbolic sigmoid.
    -------
    ave : Variable
        average kl divergence between Y and Y_hat.
    """

    assert hasattr(Y_hat, 'owner')

    owner = Y_hat.owner
    assert owner is not None
    op = owner.op

    if not hasattr(op, 'scalar_op'):
        raise ValueError("Expected Y_hat to be generated by an Elemwise "
                         "op, got "+str(op)+" of type "+str(type(op)))

    assert isinstance(op.scalar_op, TT.nnet.sigm.ScalarSigmoid)

    z, = owner.inputs

    z = z.reshape(Y.shape)
    term_1 = Y * TT.nnet.softplus(-z)
    term_2 = (1 - Y) * TT.nnet.softplus(z)

    total = term_1 + term_2
    if cost_mask is not None:
        total = cost_mask * total

    ave = total.sum() / (total.shape[1] * (total.shape[2] - 2))
    #ave = ave.mean()
    return ave


def kl(Y, Y_hat, cost_mask=None,
       batch_vec=True,
       cost_matrix=False,
       sum_tru_time=False,
       normalize_by_outsize=True):

    """
    Warning: This function expects a sigmoid nonlinearity in the
    output layer. Returns a batch (vector) of mean across units of
    KL divergence for each example,
    KL(P || Q) where P is defined by Y and Q is defined by Y_hat:
    p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)
    For binary p, some terms drop out:
    - p log q - (1-p) log (1-q)
    - p log sigmoid(z) - (1-p) log sigmoid(-z)
    p softplus(-z) + (1-p) softplus(z)
    Parameters
    ----------
    Y : Variable
        targets for the sigmoid outputs. Currently Y must be purely binary.
        If it's not, you'll still get the right gradient, but the
        value in the monitoring channel will be wrong.
    Y_hat : Variable
        predictions made by the sigmoid layer. Y_hat must be generated by
        fprop, i.e., it must be a symbolic sigmoid.
    -------
    ave : Variable
        average kl divergence between Y and Y_hat.
    """

    assert hasattr(Y_hat, 'owner')

    owner = Y_hat.owner
    assert owner is not None
    op = owner.op

    if not hasattr(op, 'scalar_op'):
        raise ValueError("Expected Y_hat to be generated by an Elemwise "
                         "op, got "+str(op)+" of type "+str(type(op)))

    assert isinstance(op.scalar_op, TT.nnet.sigm.ScalarSigmoid)

    z, = owner.inputs
    z = z.reshape(Y.shape)
    term_1 = Y * TT.nnet.softplus(-z)
    term_2 = (1 - Y) * TT.nnet.softplus(z)

    total = term_1 + term_2

    if cost_mask is not None:
        if cost_mask.ndim != total.ndim:
            cost_mask = cost_mask.dimshuffle(0, 1, 'x')

        total = cost_mask * total
    if not sum_tru_time:
        if normalize_by_outsize:
            if cost_matrix:
                ave = total.sum(-1) / TT.cast((total.shape[2] - 2), "float32")
            else:
                if batch_vec:
                    ave = total.sum(0).sum(1) / TT.cast((total.shape[2] - 2), "float32")
                else:
                    ave = total.sum() / (total.shape[1] * (total.shape[2] - 2))
        else:
            if cost_matrix:
                ave = total.sum(-1)
            else:
                if batch_vec:
                    ave = total.sum(0).sum(1)
                else:
                    ave = total.sum() / TT.cast(total.shape[1], "float32")
    else:
        assert not cost_matrix
        if normalize_by_outsize:
            if batch_vec:
                ave = cumsum(total.sum(-1) / TT.cast((total.shape[2] - 2), "float32"),
                        axis=0)[::-1]
            else:
                ave = cumsum(total.sum((1, 2)) / (total.shape[1] * (total.shape[2] - 2)),
                        axis=0)[::-1]
        else:
            if batch_vec:
                ave = cumsum(total.sum(-1), axis=0)[::-1]
            else:
                ave = cumsum(total.sum((1, 2)) / TT.cast(total.shape[1], "float32"), axis=0)[::-1]

    return ave


def _grab_probs(class_probs, target, use_fast_ver=False):
    if class_probs.ndim == 3:
        class_probs = class_probs.reshape((-1, class_probs.shape[-1]))

    shape0 = class_probs.shape[0]
    shape1 = class_probs.shape[1]

    p = None
    if target.ndim == 2 and use_fast_ver:
        target = target.flatten()
        cp = class_probs.reshape((target.shape[0], -1))
        p = TT.diag(cp.T[target])
    else:
        if target.ndim > 1:
            target = target.flatten()
        assert target.ndim == 1, 'make sure target is a vector of ints'
        assert 'int' in target.dtype
        pos = TT.arange(shape0)*shape1
        new_targ = target + pos
        p = class_probs.reshape((shape0*shape1, 1))[new_targ].reshape((shape0,))
    return p


def nll_simple(Y, Y_hat,
               cost_mask=None,
               cost_ent_mask=None,
               cost_ent_desc_mask=None):

    probs = Y_hat
    pred = TT.argmax(probs, axis=1).reshape(Y.shape)
    errors = TT.neq(pred, Y)
    ent_errors = None
    if cost_ent_mask is not None:
        pred_ent = TT.argmax(probs * cost_ent_mask.dimshuffle('x', 0),
                             axis=1).reshape(Y.shape)
        ent_errors = TT.neq(pred_ent, Y).mean()

    ent_desc_errors = None
    if cost_ent_desc_mask is not None:
        pred_desc_ent = TT.argmax(probs * cost_ent_desc_mask,
                             axis=1).reshape(Y.shape)
        ent_desc_errors = TT.neq(pred_desc_ent, Y).mean()

    LL = TT.log(_grab_probs(probs, Y) + 1e-8).reshape(Y.shape)

    if cost_mask is not None:
        total = cost_mask * LL
        errors = cost_mask * errors
        ncosts = TT.sum(cost_mask)
        mean_errors = TT.sum(errors) / (ncosts)
        ave = -TT.sum(total) / Y.shape[1]
    else:
        mean_errors = TT.mean(errors)
        ave = -TT.sum(LL) / Y.shape[0]
    return ave, mean_errors, ent_errors, ent_desc_errors


def nll(Y, Y_hat, cost_mask=None, batch_vec=True,
        cost_matrix=False, use_fast_ver=False):
    probs = Y_hat
    pred = TT.argmax(probs, axis=-1).reshape(Y.shape)
    errors = TT.neq(pred, Y)
    LL = TT.log(_grab_probs(probs, Y, use_fast_ver=use_fast_ver) + 1e-8).reshape(Y.shape)

    if cost_mask is not None:
        total = cost_mask * LL
        errors = cost_mask * errors
        ncosts = TT.sum(cost_mask)
        mean_errors = TT.sum(errors) / (ncosts)
        if cost_matrix:
            ave = -total
        else:
            if batch_vec:
                ave = -total.sum(0)
            else:
                ave = -TT.sum(total) / Y.shape[1]
    else:
        mean_errors = TT.mean(errors)
        if cost_matrix:
            ave = -LL
        else:
            if batch_vec:
               ave = -LL.sum(0)
            else:
               ave = -TT.sum(LL) / Y.shape[1]

    return ave, mean_errors


def multiclass_hinge_loss(predictions, targets, delta=1):
    """Computes the multi-class hinge loss between predictions and targets.
    .. math:: L_i = \\max_{j \\not = p_i} (0, t_j - t_{p_i} + \\delta)
    Parameters
    ----------
    predictions : Theano 2D tensor
        Predictions in (0, 1), such as softmax output of a neural network,
        with data points in rows and class probabilities in columns.
    targets : Theano 2D tensor or 1D tensor
        Either a vector of int giving the correct class index per data point
        or a 2D tensor of one-hot encoding of the correct class in the same
        layout as predictions (non-binary targets in [0, 1] do not work!)
    delta : scalar, default 1
        The hinge loss margin
    Returns
    -------
    Theano 1D tensor
        An expression for the item-wise multi-class hinge loss
    Notes
    -----
    This is an alternative to the categorical cross-entropy loss for
    multi-class classification problems
    """
    num_cls = predictions.shape[1]
    if targets.ndim == predictions.ndim - 1:
        targets = theano.tensor.extra_ops.to_one_hot(targets, num_cls)
    elif targets.ndim != predictions.ndim:
        raise TypeError('rank mismatch between targets and predictions')
    corrects = predictions[targets.nonzero()]
    rest = theano.tensor.reshape(predictions[(1-targets).nonzero()],
                                 (-1, num_cls-1))
    rest = theano.tensor.max(rest, axis=1)
    return theano.tensor.nnet.relu(rest - corrects + delta).mean(0).sum()
Warning: This function expects a sigmoid nonlinearity in the     output layer. Returns a batch (vector) of mean across units of     KL divergence for each example,     KL(P || Q) where P is defined by Y and Q is defined by Y_hat:     p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)     For binary p, some terms drop out:     - p log q - (1-p) log (1-q)     - p log sigmoid(z) - (1-p) log sigmoid(-z)     p softplus(-z) + (1-p) softplus(z)     Parameters     ----------     Y : Variable         targets for the sigmoid outputs. Currently Y must be purely binary.         If it's not, you'll still get the right gradient, but the         value in the monitoring channel will be wrong.     Y_hat : Variable         predictions made by the sigmoid layer. Y_hat must be generated by         fprop, i.e., it must be a symbolic sigmoid.     -------     ave : Variable         average kl divergence between Y and Y_hat. #ave = ave.mean() Warning: This function expects a sigmoid nonlinearity in the     output layer. Returns a batch (vector) of mean across units of     KL divergence for each example,     KL(P || Q) where P is defined by Y and Q is defined by Y_hat:     p log p - p log q + (1-p) log (1-p) - (1-p) log (1-q)     For binary p, some terms drop out:     - p log q - (1-p) log (1-q)     - p log sigmoid(z) - (1-p) log sigmoid(-z)     p softplus(-z) + (1-p) softplus(z)     Parameters     ----------     Y : Variable         targets for the sigmoid outputs. Currently Y must be purely binary.         If it's not, you'll still get the right gradient, but the         value in the monitoring channel will be wrong.     Y_hat : Variable         predictions made by the sigmoid layer. Y_hat must be generated by         fprop, i.e., it must be a symbolic sigmoid.     -------     ave : Variable         average kl divergence between Y and Y_hat. Computes the multi-class hinge loss between predictions and targets.     .. math:: L_i = \\max_{j \\not = p_i} (0, t_j - t_{p_i} + \\delta)     Parameters     ----------     predictions : Theano 2D tensor         Predictions in (0, 1), such as softmax output of a neural network,         with data points in rows and class probabilities in columns.     targets : Theano 2D tensor or 1D tensor         Either a vector of int giving the correct class index per data point         or a 2D tensor of one-hot encoding of the correct class in the same         layout as predictions (non-binary targets in [0, 1] do not work!)     delta : scalar, default 1         The hinge loss margin     Returns     -------     Theano 1D tensor         An expression for the item-wise multi-class hinge loss     Notes     -----     This is an alternative to the categorical cross-entropy loss for     multi-class classification problems