import os
import random
import time

import cv2
import numpy as np
import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import save_image

from datasets import get_subclass_dataset


def set_random_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)


IMAGENET_PATH = "~/data/ImageNet"


check = time.time()

transform = transforms.Compose(
    [
        transforms.Resize(256),
        transforms.CenterCrop(256),
        transforms.Resize(32),
        transforms.ToTensor(),
    ]
)

# remove airliner(1), ambulance(2), parking_meter(18), schooner(22) since similar class exist in CIFAR-10
class_idx_list = list(range(30))
remove_idx_list = [1, 2, 18, 22]
for remove_idx in remove_idx_list:
    class_idx_list.remove(remove_idx)

set_random_seed(0)
train_dir = os.path.join(IMAGENET_PATH, "one_class_train")
Imagenet_set = datasets.ImageFolder(train_dir, transform=transform)
Imagenet_set = get_subclass_dataset(Imagenet_set, class_idx_list)
Imagenet_dataloader = DataLoader(
    Imagenet_set, batch_size=100, shuffle=True, pin_memory=False
)

total_test_image = None
for n, (test_image, target) in enumerate(Imagenet_dataloader):

    if n == 0:
        total_test_image = test_image
    else:
        total_test_image = torch.cat((total_test_image, test_image), dim=0).cpu()

    if total_test_image.size(0) >= 10000:
        break

print(f"Preprocessing time {time.time()-check}")

if not os.path.exists("./Imagenet_fix"):
    os.mkdir("./Imagenet_fix")

check = time.time()
for i in range(10000):
    save_image(total_test_image[i], f"Imagenet_fix/correct_resize_{i}.png")
print(f"Saving time {time.time()-check}")
