def restore_rgb_img(batch_img, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]):
    """
    restore the rgb img
    Args:
        batch_img: rgb image with normalization, shape= (batch_size, num_frames, C, H, W)
        mean:      params to restore the rgb image
        std:       params to restore the rgb image

    Returns:

    """
 

    assert len(batch_img.shape) == 5
    assert batch_img.shape[2] == 3
    C = batch_img.shape[2]
    for channel_index, channle_mean, channel_std in zip(range(C), mean, std):
        batch_img[:, :, [channel_index], :, :] = batch_img[:, :, [channel_index], :, :]*channel_std + channle_mean
    return
 
 batch_img


def vis_cos_dist(cos_dist, batch_idx,  width, height, epoch):  # batch_idx.shape (16, 3, 4, 256, 128) (b, c, t, h ,w)
    GRID_SPACING = 10
    imagenet_mean = [0.485, 0.456, 0.406]
    imagenet_std = [0.229, 0.224, 0.225]
    if not os.path.exists('./similarity_heat_map'):
        os.makedirs('./similarity_heat_map'
 
)
    imgs = restore_rgb_img(batch_idx.permute(0, 2, 1, 3, 4))
    # imgs = (imgs**2).sum(1)
    imgs = imgs.cpu()

    for i in range(len(cos_dist)):
        for j in range(len(cos_dist[i])):
            for target in list(cos_dist[i][j]):  #  target = 'ap_in'  'ap_out'
                if cos_dist[i][j][target] is not None:

                    outputs = cos_dist[i][j][target]  # outputs is a similarity matrix  (b, t, h ,w)
                    outputs = (outputs**2).sum(1)
                    
                    b, h ,w = outputs.size()
                    outputs = outputs.view(b, h*w)
                    outputs = F.normalize(outputs, p=2, dim=1)
                    outputs = outputs.view(b, h, w)
                    outputs = outputs.cpu()

                    for k in range(outputs.size(0)):

                        img = imgs[k,2, ...]
                        for t, m, s in zip(img, imagenet_mean, imagenet_std):
                            t.mul_(s).add_(m).clamp_(0, 1)
                        img_np = np.uint8(np.floor(img.numpy() * 255))
                        img_np = img_np.transpose((1, 2, 0)) # (c, h, w) -> (h, w, c)

                        # activation map
                        am = outputs[k, ...].detach().numpy()
                        am = cv2.resize(am, (width, height))
                        am = 255 * (am - np.max(am)) / (np.max(am) - np.min(am) + 1e-12)
                        am = np.uint8(np.floor(am))
                        am = cv2.applyColorMap(am, cv2.COLORMAP_JET)
                        
                        # overlapped
                        overlapped = img_np * 0.3 + am * 0.7
                        overlapped[overlapped>255] = 255
                        overlapped = overlapped.astype(np.uint8)

                        # save images in a single figure (add white spacing between images)
                        # from left to right: original image, activation map, overlapped image
                        grid_img = 255 * np.ones((height, 3*width+2*GRID_SPACING, 3), dtype=np.uint8)
                        grid_img[:, :width, :] = img_np[:, :, ::-1]
                        grid_img[:, width+GRID_SPACING: 2*width+GRID_SPACING, :] = am
                        grid_img[:, 2*width+2*GRID_SPACING:, :] = overlapped
                        
                        path = './similarity_heat_map/epoch{}_batch{}_layer{}_block{}_{}.jpg'.format(epoch, k, i, j, target)
                        cv2.imwrite(path, grid_img)