import torch


def Hess_img_distmat(ImDist, img_all, nrow=11):
    """
    img_all: a torch 4D array of shape [N, 3, 256, 256] on cpu()
    nrow: specify how many images are generated in one axis.
        It will arrange the images in `img_all` in a matrix and calculate the distance to the center image in each row.
    """
    distmat = torch.zeros(img_all.shape[0]).view(-1,nrow)
    nTot = distmat.shape[0]
    for irow in range(nTot):
        rstr = irow * nrow
        rend = nrow + rstr
        rmid = (nrow-1)//2 + rstr
        with torch.no_grad():
            dists = ImDist(img_all[rstr:rend,:],img_all[rmid,:]).squeeze()
        distmat[irow, :]=dists.cpu()
    return distmat