import torch


def cox_loss_ties(pred, cens, tril, tied_matrix):
    """
    Compute the Efron version of the Cox loss. This version take into
    account the ties.
    t unique time
    H_t denote the set of indices i such that y^i = t and c^i =1.
    c^i = 1 event occured.
    m_t = |H_t| number of elements in H_t.
    l(theta) = sum_t (sum_{i in H_t} h_{theta}(x^i)
                     - sum_{l=0}^{m_t-1} log (
                        sum_{i: y^i >= t} exp(h_{theta}(x^i))
                        - l/m_t sum_{i in H_t} exp(h_{theta}(x^i)))

    Parameters
    ----------
    pred : torch tensor
        Model prediction.
    cens : torch tensor
        Event tensor.
    tril : torch tensor
        Lower triangular tensor.
    tied_matrix : torch tensor
        Diagonal by block tensor.

    Returns
    -------
    loss : float
        Efron version of the Cox loss.
    """

    # Note that the observed variable is not required as we are sorting the
    # inputs when generating the batch according to survival time.

    # exp(h_{theta}(x^i))
    exp_pred = torch.exp(pred)
    # Term corresponding to the sum over events in the risk pool
    # sum_{i: y^i >= t} exp(h_{theta}(x^i))
    future_theta = torch.mm(tril.transpose(1, 0), exp_pred)
    # sum_{i: y^i >= t} exp(h_{theta}(x^i))
    # - l/m_t sum_{i in H_t} exp(h_{theta}(x^i))
    tied_term = future_theta - torch.mm(tied_matrix, exp_pred)
    # log (sum_{i: y^i >= t} exp(h_{theta}(x^i))
    #      - l/m_t sum_{i in H_t} exp(h_{theta}(x^i))
    tied_term = torch.log(tied_term)
    # event row vector to column
    tied_term = tied_term.view((-1, 1))
    cens = cens.view((-1, 1))
    # sum_t (sum_{i in H_t} h_{theta}(x^i)
    #       - sum_{l=0}^{m_t-1} log (
    #          sum_{i: y^i >= t} exp(h_{theta}(x^i))
    #          - l/m_t sum_{i in H_t} exp(h_{theta}(x^i)))
    loss = (pred - tied_term) * cens
    # Negative loglikelihood
    loss = -torch.mean(loss)
    return loss