import torch
import os
from torchvision import transforms
from PIL import Image


class VehicleIdDataset(torch.utils.data.Dataset):
    def __init__(self, ds_path, mode, test_set_type=None, transforms_mode='test',
                 train_start_class=0, train_fin_class=13164,
                 test_set_start_class=13164, small_test_set_fin_class=13964,
                 medium_test_set_fin_class=14764, large_test_set_fin_class=15564):
        self.ds_path = ds_path
        self.mode = mode
        self.test_set_type = test_set_type
        self.transforms_mode = transforms_mode
        self.sz_crop = 224
        self.sz_resize = 256
        self.mean = [0.485, 0.456, 0.406]
        self.std = [0.229, 0.224, 0.225]

        if self.mode == 'train':
            self.classes = range(train_start_class, train_fin_class)
            file_name = 'train_list.txt'
        else:
            if self.test_set_type is 'small':
                test_set_fin_class = small_test_set_fin_class
                file_name = 'test_list_800.txt'
            elif self.test_set_type is 'medium':
                test_set_fin_class = medium_test_set_fin_class
                file_name = 'test_list_1600.txt'
            else:
                test_set_fin_class = large_test_set_fin_class
                file_name = 'test_list_2400.txt'
            self.classes = range(test_set_start_class, test_set_fin_class)

        self.img_paths, self.labels, self.idxs = [], [], []
        im_paths_list = os.path.join(ds_path, f'train_test_split/{file_name}')
        with open(im_paths_list) as f:
            t = [t.replace('\n', '').split() for t in f.readlines()]
        for i, t_ in enumerate(t):
            self.img_paths.append(ds_path + '/image/' + t_[0] + '.jpg')
            self.labels.append(int(t_[1]))
            self.idxs.append(i)
        assert len(set(self.labels)) == len(set(self.classes))

    def __len__(self):
        return len(self.labels)

    def __getitem__(self, index):
        im = Image.open(self.img_paths[index])
        if len(list(im.split())) == 1:
            im = im.convert('RGB')
        im = self.apply_augmentation(im)
        return im, self.labels[index], index

    def apply_augmentation(self, im):
        if self.transforms_mode == 'train':
            transforms_ = transforms.Compose([
                transforms.RandomResizedCrop(self.sz_crop),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=self.mean,
                    std=self.std,
                )
            ])
        else:
            transforms_ = transforms.Compose([
                transforms.Resize(self.sz_resize),
                transforms.CenterCrop(self.sz_crop),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=self.mean,
                    std=self.std,
                )
            ])
        return transforms_(im)

    def get_nb_classes(self):
        return len(self.classes)

    def get_label(self, index):
        return self.labels[index]

    def get_within_indexes(self, indexes):
        self.labels = [self.labels[i] for i in indexes]
        self.idxs = [self.idxs[i] for i in indexes]
        self.img_paths = [self.img_paths[i] for i in indexes]
