import torch


def uc_jsd(nll_valid, nll_train, nll_sampled):
    """ Computes the UC-JSD (metric used for the benchmark of generative models
    in Arús-Pous, J. et al., J. Chem. Inf., 2019, 1-13).

    Args:
      nll_valid (torch.Tensor) : Contains NLLs for sampling the correct action
        of structures in the validation set.
      nll_train (torch.Tensor) : Contains NLLs for sampling the correct action
        of structures in the training set.
      nll_sampled (torch.Tensor) : Contains NLLs for sampling the correct action
        of structures in the generated set.
    """
    min_len = min(len(nll_valid), len(nll_sampled), len(nll_train))

    # make all the distributions the same length (dim=0)
    nll_valid_norm = nll_valid[:min_len] / torch.sum(nll_valid[:min_len])
    nll_train_norm = nll_train[:min_len] / torch.sum(nll_train[:min_len])
    nll_sampled_norm = nll_sampled[:min_len] / torch.sum(nll_sampled[:min_len])

    nll_sum = (nll_valid_norm + nll_train_norm + nll_sampled_norm) / 3

    uc_jsd = (torch.nn.functional.kl_div(nll_valid_norm, nll_sum)
              + torch.nn.functional.kl_div(nll_train_norm, nll_sum)
              + torch.nn.functional.kl_div(nll_sampled_norm, nll_sum)) / 3

    return float(uc_jsd)