import os
import torch
import glob
import random
import os
from torch.utils.data import Dataset
from PIL import Image
import torchvision
from torchvision.utils import make_grid
import matplotlib.pyplot as plt



def show_tensor_images(image_tensor, num_images=25, size=(1, 28, 28)):
    image_tensor = (image_tensor + 1) / 2
    image_shifted = image_tensor
    image_unflat = image_shifted.detach().cpu().view(-1, *size)
    image_grid = make_grid(image_unflat[:num_images], nrow=5)
    plt.imshow(image_grid.permute(1, 2, 0).squeeze())
    plt.show()


class ImageDataset(Dataset):
    def __init__(self, root, transform=None, mode="train"):
        self.transform = transform
        self.files_A = sorted(glob.glob(os.path.join(root, "%sA" % mode) + "/*.*"))
        self.files_B = sorted(glob.glob(os.path.join(root, "%sB" % mode) + "/*.*"))
        if len(self.files_A) > len(self.files_B):
            self.files_A, self.files_B = self.files_B, self.files_A
        self.new_perm()
        assert(len(self.files_A) > 0, "Make sure to have downloaded the horze2zebra images")

    def new_perm(self):
        self.randperm = torch.randperm(len(self.files_B))[:len(self.files_A)]

    def __getitem__(self, index):
        item_A = self.transform(Image.open(self.files_A[index % len(self.files_A)]))
        item_B = self.transform(Image.open(self.files_B[self.randperm[index]]))

        if item_A.shape[0] != 3:
            item_A = item_A.repeat(3, 1, 1)
        if item_B.shape[0] != 3:
            item_B = item_B.repeat(3, 1, 1)
        if index == len(self) - 1:
            self.new_perm()
        # Normalization
        return (item_A - 0.5) * 2, (item_B - 0.5) * 2
    
    def __len__(self):
        return min(len(self.files_A), len(self.files_B))


def get_last_training_step(weight_paths):
    weights_list = os.listdir("/content/drive/MyDrive/CycleGAN_weights/")
    max_step = 0
    for w in weights_list:
        step = w.split("_")[-1]
        step = int(step.split(".")[0])
        if step > max_step:
            max_step = step
    return max_step
