import torch
import numpy as np
import matplotlib.pyplot as plt

import torchvision.transforms.functional as F

plt.rcParams["savefig.bbox"] = 'tight'

def show(imgs, save_fig_path='grid_num.png'):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fix, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

    plt.savefig(save_fig_path)

def view_predict(vae, n=30, figsize=15, save_fig_path='vae_predict.png'):
    digit_size = 28
    scale = 1.0
    figure = np.zeros((digit_size * n, digit_size * n))
    
    grid_x = np.linspace(-scale, scale, n)
    grid_y = np.linspace(-scale, scale, n)[::-1]

    vae.eval()
    vae.cpu()
    for i, yi in enumerate(grid_y):
        for j, xi in enumerate(grid_x):
            z_sample = torch.tensor([[xi, yi]]).float().cpu()
            y = vae.decoder(z_sample)
            digit = y[0].view(digit_size, digit_size).cpu().detach().numpy()
            figure[
                i * digit_size : (i + 1) * digit_size,
                j * digit_size : (j + 1) * digit_size,
            ] = digit

    plt.figure(figsize=(figsize, figsize))
    
    start_range = digit_size // 2
    end_range = n * digit_size + start_range
    pixel_range = np.arange(start_range, end_range, digit_size)
    
    sample_range_x = np.round(grid_x, 1)
    sample_range_y = np.round(grid_y, 1)
    
    plt.xticks(pixel_range, sample_range_x)
    plt.yticks(pixel_range, sample_range_y)
    plt.xlabel("z[0]")
    plt.ylabel("z[1]")
    plt.imshow(figure, cmap="Greys_r")
    plt.savefig(save_fig_path)
    plt.close()

def view_label_clusters(vae, loader, save_fig_path='vae_clusters.png'):
    vae.eval()
    vae.cpu()

    z_mean = []
    labels = []
    for img, label in loader:
        _, mean, _ = vae.encoder(img, x_device=torch.device('cpu'))
        z_mean.extend(mean.detach().cpu().numpy())
        labels.extend(label.detach().cpu().numpy())

    z_mean = np.array(z_mean)
    labels = np.array(labels)

    print(z_mean.shape, labels.shape)

    plt.figure(figsize=(12, 10))
    plt.scatter(z_mean[:, 0], z_mean[:, 1], c=labels)
    plt.colorbar()
    plt.xlabel("z[0]")
    plt.ylabel("z[1]")
    plt.savefig(save_fig_path)
    plt.close()


