import torch


def ume_ustat_h1_mean_variance(feature_matrix, return_variance=True, 
                               use_unbiased=True):
    """
    Compute the mean and variance of the asymptotic normal distribution
    under H1 of the test statistic. The mean converges to a constant as
    n->\infty.

    feature_matrix: n x J feature matrix
    return_variance: If false, avoid computing and returning the variance.
    use_unbiased: If True, use the unbiased version of the mean. Can be
        negative.

    Return the mean [and the variance]
    """
    Z = feature_matrix
    n = Z.size(0)
    assert n > 1, 'Need n > 1 to compute the mean of the statistic.'
    if use_unbiased:
        # t1 = np.sum(np.mean(Z, axis=0)**2)*(n/float(n-1))
        t1 = torch.sum(torch.mean(Z, dim=0)**2).mul(n).div(n-1)
        # t2 = np.mean(np.sum(Z**2, axis=1))/float(n-1)
        t2 = torch.mean(torch.sum(Z**2, dim=1)).div(n-1)
        mean_h1 = t1 - t2
    else:
        # mean_h1 = np.sum(np.mean(Z, axis=0)**2)
        mean_h1 = torch.sum(torch.mean(Z, dim=0)**2)

    if return_variance:
        # compute the variance 
        # mu = np.mean(Z, axis=0)  # length-J vector
        mu = torch.mean(Z, dim=0)  # length-J vector
        # variance = 4.0*np.mean(np.dot(Z, mu)**2) - 4.0*np.sum(mu**2)**2
        variance = 4.0*torch.mean(torch.matmul(Z, mu)**2) - 4.0*torch.sum(mu**2)**2
        return mean_h1, variance
    else:
        return mean_h1