import torch


def compute_loss(X_obs, Y_obs, Y_obs_bj, n_obs_ot, batch_size, eps=1e-10,
                 weight=0.5, M_obs=None):
    """
    loss function from the paper
    :param X_obs: torch.tensor, the true X values at the observations
    :param Y_obs: torch.tensor, the predicted values at the observation
    :param Y_obs_bj: torch.tensor, the predicted values before the jump at the
            observation
    :param n_obs_ot: torch.tensor, the number of observations over the entire
            time-line for each element of the batch
    :param batch_size: int or float
    :param eps: float, a small constant which is added before taking torch.sqrt
            s.t. the sqrt never becomes zero (which would yield NaNs for the
            gradient)
    :param weight: float in [0,1], weighting of the two parts of the loss
            function,
                0.5: standard loss as described in paper
                (0.5, 1): more weight to be correct after the jump, can be
                theoretically justified similar to standard loss
                1: only the value after the jump is trained
    :param M_obs: None or torch.tensor, if not None: same size as X_obs with
            0  and 1 entries, telling which coordinates were observed
    :return: torch.tensor (with the loss, reduced to 1 dim)
    """
    if M_obs is None:
        inner = (2*weight * torch.sqrt(torch.sum((X_obs - Y_obs) ** 2, dim=1)+eps) +
                 2*(1-weight) * torch.sqrt(torch.sum((Y_obs_bj - Y_obs) ** 2, dim=1)
                                           + eps)) ** 2
    else:
        inner = (2 * weight * torch.sqrt(
            torch.sum(M_obs * (X_obs - Y_obs) ** 2, dim=1) + eps) +
                 2 * (1 - weight) * torch.sqrt(
                    torch.sum(M_obs * (Y_obs_bj - Y_obs) ** 2, dim=1)
                    + eps)) ** 2
    outer = torch.sum(inner / n_obs_ot)
    return outer / batch_size