import os
import sys  
import torch
import random
import numpy as np
import PIL.Image as PImage
import torchvision.datasets as datasets
import torch.utils.data as data
from PIL import Image, ImageOps, ImageFilter
import torchvision
from torchvision.datasets import ImageFolder
from torchvision import transforms
from torch.utils.data import Dataset, DataLoader
from torch.utils.data.distributed import DistributedSampler
from data.augmentation import random_crop_arr, center_crop_arr

paths = {
    "CIFAR10": "CIFAR10",
    "SVHN": "SVHN",
}

def build_transform():
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
    ])
    return transform

def build_dataloader(args):
    data_path = os.path.join(args.dataset_dir, paths[args.dataset_name])
    transform = build_transform()

    if args.dataset_name == "CIFAR10":
        train_set = torchvision.datasets.CIFAR10(root=data_path, train=True, transform=transform, download=True)
        val_set = torchvision.datasets.CIFAR10(root=data_path, train=False, transform=transform, download=True)
    elif args.dataset_name == "SVHN":
        train_set = torchvision.datasets.SVHN(root=data_path, split='train', transform=transform, download=True)
        val_set = torchvision.datasets.SVHN(root=data_path, split='test', transform=transform, download=True)

    print("dataset name:", args.dataset_name)
    print("len train_set:", len(train_set))
    print("len val_set:", len(val_set))

    #train_sampler = torch.utils.data.distributed.DistributedSampler(train_set)
    train_dataloader = DataLoader(
        dataset=train_set, num_workers=args.workers, pin_memory=True,
        batch_size=args.batch_size, shuffle=True, drop_last=True
    )

    val_dataloader = DataLoader(
        dataset=val_set, num_workers=args.workers, pin_memory=True,
        batch_size=args.batch_size*2, shuffle=False, drop_last=False
    )
    print("len train_dataloader:", len(train_dataloader))
    print("len val_dataloader:", len(val_dataloader))
    return train_dataloader, val_dataloader, len(train_set), len(val_set)