import torch


def init_center_c(train_loader, dataset_name, net, device, eps=0.1,):
    """Initialize hypersphere center c as the mean from an initial forward pass on the data."""
    n_samples = 0
    
    if dataset_name == "mnist":
        tnum = 32
    elif dataset_name == "cifar10":
        tnum = 96*2*2
    elif dataset_name == "ivc-filter" or dataset_name == "rsna":
        tnum = 512
    else:
        tnum = 0
    c = torch.zeros(tnum, device=device)

    net.eval()
    with torch.no_grad():
        for idx, inputs in enumerate(train_loader):
            # get the inputs of the batch
            data, _ = inputs
            org_imgs, _ = data
            org_imgs = org_imgs.to(device)
            outputs = net.get_embedding(org_imgs)
            n_samples += outputs.shape[0]
            c += torch.sum(outputs, dim=0)

    c /= n_samples

    # If c_i is too close to 0, set to +-eps. Reason: a zero unit can be trivially matched with zero weights.
    c[(abs(c) < eps) & (c < 0)] = -eps
    c[(abs(c) < eps) & (c > 0)] = eps

    return c