# A file of bregman divergences
import torch


def mse(x, y):
    sq_diff = (x - y) ** 2
    if x.shape != y.shape:
        assert False, "x and y must have the same shape"
    return sq_diff.reshape(sq_diff.size(0), -1).sum(dim=-1)


# TODO: check if this formulation is correct
def jump_kernel_elbo(x, y, eps=1e-6):
    # x_safe: true length
    # y_safe: predicted length
    x_safe = torch.clamp(x, min=eps)
    y_safe = torch.clamp(y, min=eps)

    return y_safe - x_safe + x_safe * (torch.log(x_safe) - torch.log(y_safe))