import sys
import os.path as osp
from torchvision.datasets import STL10

from dassl.utils import mkdir_if_missing


def extract_and_save_image(dataset, save_dir):
    if osp.exists(save_dir):
        print('Folder "{}" already exists'.format(save_dir))
        return

    print('Extracting images to "{}" ...'.format(save_dir))
    mkdir_if_missing(save_dir)

    for i in range(len(dataset)):
        img, label = dataset[i]
        if label == -1:
            label_name = "none"
        else:
            label_name = str(label)
        imname = str(i).zfill(6) + "_" + label_name + ".jpg"
        impath = osp.join(save_dir, imname)
        img.save(impath)


def download_and_prepare(root):
    train = STL10(root, split="train", download=True)
    test = STL10(root, split="test")
    unlabeled = STL10(root, split="unlabeled")

    train_dir = osp.join(root, "train")
    test_dir = osp.join(root, "test")
    unlabeled_dir = osp.join(root, "unlabeled")

    extract_and_save_image(train, train_dir)
    extract_and_save_image(test, test_dir)
    extract_and_save_image(unlabeled, unlabeled_dir)


if __name__ == "__main__":
    download_and_prepare(sys.argv[1])