import os
import matplotlib.pyplot as plt
from torchvision.utils import make_grid
import torch

def show_tensor_images(image_tensor, epoch,step,otype,num_images=25, size=(1, 28, 28)):
    image_shifted = image_tensor
    image_unflat = image_shifted.detach().cpu().view(-1, *size)
    image_grid = make_grid(image_unflat[:num_images], nrow=5)
    if not os.path.exists(f"/outputs/Epoch{epoch}"):
        os.makedirs(f"/outputs/Epoch{epoch}")
    plt.savefig(os.path.join(f"/outputs/Epoch{epoch}_step_{step}_{otype}"))
    plt.close()


def crop(image, new_shape):
    middle_height = image.shape[2] // 2
    middle_width = image.shape[3] // 2
    starting_height = middle_height - new_shape[2] // 2
    final_height = starting_height + new_shape[2]
    starting_width = middle_width - new_shape[3] // 2
    final_width = starting_width + new_shape[3]
    cropped_image = image[:, :, starting_height:final_height, starting_width:final_width]    
    return cropped_image


def get_gen_loss(gen, disc, real, condition, dev_criterion, recon_criterion, lambda_recon):
    gen_pred = gen(condition)
    disc_pred = disc(gen_pred, real)
    gen_adv_loss = dev_criterion(disc_pred, torch.ones_like(disc_pred))
    gen_recon_loss = recon_criterion(real,gen_pred)
    gen_loss = gen_adv_loss + lambda_recon * gen_recon_loss

    return gen_loss