import numpy as np
from matplotlib import pyplot as plt

import torch
from torch import nn

def get_heatmap(patch_scores, fmap_dims, patch_size, patch_stride, img_size):
    """ 
        Returns a heatmap of *img_size*
        *patch_scores* is either patch pred probabilities or attention weights
    """

    device = patch_scores.device
    num_patches_y = int((fmap_dims[1] - patch_size) / patch_stride + 1)

    b, k = patch_scores.shape
    heatmap = torch.zeros((b, 1, fmap_dims[0], fmap_dims[1])).to(device)
    heatmap_norm  = torch.zeros_like(heatmap)
    for i in range(b):
        for j in range(k):
            patch_score = patch_scores[i,j].squeeze()
            row_id = int(j // num_patches_y)
            col_id = int(j % num_patches_y)
            heatmap[i, 0, row_id : row_id + patch_size, col_id : col_id + patch_size] += patch_score
            heatmap_norm[i, 0, row_id : row_id + patch_size, col_id : col_id + patch_size] += torch.ones_like(patch_score)
    heatmap /= heatmap_norm
    heatmap = nn.functional.interpolate(heatmap, (img_size, img_size), mode = 'bilinear')
    return heatmap

def vis_multiple_imgs(images, heatmaps, fpath):
    """ 
        Saves multiple *images* with *heatmaps* overlays in *fpath*
    """

    b = images.shape[0]
    num_rows = int(b**0.5)
    num_cols = num_rows
    
    fig = plt.figure(figsize=(8,8), dpi=600)
    for i, (img, hmap) in enumerate(zip(images, heatmaps)):
        img = img.cpu().numpy()
        img = (img - np.min(img)) / np.ptp(img)
        img = np.transpose(img, (1, 2, 0))

        hmap = hmap.cpu().numpy()
        hmap = np.transpose(hmap, (1, 2, 0))

        ax = fig.add_subplot(num_rows, num_cols, i + 1)
        im_img = ax.imshow(img, alpha = 1., interpolation = 'bilinear', cmap = 'gray')
        im_hmap = ax.imshow(hmap, alpha = 0.5, interpolation = 'bilinear', cmap = 'viridis', vmin = 0., vmax = 1.)            
    
        ax.tick_params(
            axis = 'both',       
            which = 'both',      
            bottom = False,      
            top = False,         
            left = False,
            labelbottom = False,
            labelleft = False
        ) 
        
    plt.subplots_adjust(wspace = 0, hspace = 0)
    plt.savefig(fpath, bbox_inches = 'tight', pad_inches = 0, format = 'png', transparent = False)