import torch


def hsic(k_x: torch.Tensor, k_y: torch.Tensor, centered: bool = False, unbiased: bool = True) -> torch.Tensor:
    """Compute Hilbert-Schmidt Independence Criteron (HSIC)

    :param k_x: n by n values of kernel applied to all pairs of x data
    :param k_y: n by n values of kernel on y data
    :param centered: whether or not at least one kernel is already centered
    :param unbiased: if True, use unbiased HSIC estimator of Song et al (2007), else use original estimator of Gretton et al (2005)
    :return: scalar score in [0*, inf) measuring dependence of x and y

    * note that if unbiased=True, it is possible to get small values below 0.
    """
    if k_x.size() != k_y.size():
        raise ValueError("RDMs must have the same size!")
    n = k_x.size()[0]

    if not centered:
        h = torch.eye(n, device=k_y.device, dtype=k_y.dtype) - 1/n
        k_y = h @ k_y @ h

    if unbiased:
        # Remove the diagonal
        k_x = k_x * (1 - torch.eye(n, device=k_x.device, dtype=k_x.dtype))
        k_y = k_y * (1 - torch.eye(n, device=k_y.device, dtype=k_y.dtype))
        # Equation (4) from Song et al (2007)
        return ((k_x *k_y).sum() - 2*(k_x.sum(dim=0)*k_y.sum(dim=0)).sum()/(n-2) + k_x.sum()*k_y.sum()/((n-1)*(n-2))) / (n*(n-3))
    else:
        # The original estimator from Gretton et al (2005)
        return torch.sum(k_x * k_y) / (n - 1)**2