import torch
from torch.utils.data import Dataset
import json
import os
from PIL import Image
import numpy as np
from torchvision import transforms
import yaml


class IHDataset(Dataset):
    """Intracranial Hemorrhage dataset."""
    def __init__(self, root_dir, stage='train', transform=None):
        self.root_dir = root_dir
        self.stage = stage
        self.transform = transform
        self.data_dir = os.path.join(self.root_dir, self.stage)

        with open(os.path.join(root_dir, f'annots/{stage}.csv'), 'r') as csv:
            lines = [line.strip().split(',') for line in csv.readlines()]
            header, lines = lines[0], lines[1:]
            # k: dicom ID, v: IH/noIH (1/0)
            self.annots = {k:[] for k in header}
            for line in lines:
                for k,v in zip(header,line):
                    self.annots[k].append(v)

        self.idx_to_slice_id = dict(enumerate(self.annots['ID']))
        self.slice_id_to_idx = {v:k for k,v in self.idx_to_slice_id.items()}

        test_patients_path = os.path.join(root_dir,
                                          f'annots/patient_stats_test.yaml')
        self.test_patients = yaml.safe_load(open(test_patients_path, 'r'))

    def __len__(self):
        return len(self.annots['ID'])

    def __getitem__(self, idx):
        img_name = f'{self.idx_to_slice_id[idx]}.png'
        img_path = os.path.join(self.data_dir, img_name)
        sample = Image.open(img_path)

        if self.transform:
            sample = self.transform(sample)

        target = int(self.annots['IH'][idx])
        return sample, target

    def getSlice(self, slice_id):
        idx = self.slice_id_to_idx[slice_id]
        return self.__getitem__(idx)


class IHTestDataset(Dataset):
    """Intracranial Hemorrhage dataset."""
    def __init__(self, root_dir, stage='', transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.data_dir = root_dir if root_dir[-1] == '/' else root_dir + '/'
        import glob
        self.data = glob.glob(self.data_dir + '*')
        self.idx_to_img_path = dict(enumerate(self.data))

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

    def __getitem__(self, idx):
        img_path = self.idx_to_img_path[idx]
        sample = Image.open(img_path)
        if self.transform:
            sample = self.transform(sample)
        return sample, img_path


def get_datasets(conf):
    dataset = conf['data']['name']
    root_dir = conf['data']['path']
    train_transform = conf['train_transform']
    valid_transform = conf['valid_transform']
    test_transform = conf['test_transform']
    if dataset == 'IHDataset':
        train_dataset = IHDataset(root_dir=root_dir,
                                  stage='train',
                                  transform=train_transform)
        valid_dataset = IHDataset(root_dir=root_dir,
                                  stage='valid',
                                  transform=valid_transform)
        test_dataset = IHDataset(root_dir=root_dir,
                                  stage='test',
                                  transform=test_transform)
        patients_dataset = IHDataset(root_dir=root_dir,
                                     stage='test_no_balanced',
                                     transform=test_transform)
    elif dataset == 'IHTestDataset':
        return None, None, IHTestDataset(root_dir=root_dir,
                                         transform=test_transform), None
    else:
        print('Dataset {dataset} not supported.')
        exit()
    return train_dataset, valid_dataset, test_dataset, patients_dataset


def get_dataloaders(conf):
    from torch.utils.data import DataLoader
    train_dataset = conf['train_dataset']
    valid_dataset = conf['valid_dataset']
    test_dataset = conf['test_dataset']
    batch_size = conf['data']['batch_size']
    num_workers = conf['data']['num_workers']
    train_loader = DataLoader(dataset=train_dataset,
                              batch_size=batch_size,
                              num_workers=num_workers,
                              shuffle=True,
                              pin_memory=True)
    valid_loader = DataLoader(dataset=valid_dataset,
                              batch_size= batch_size,
                              num_workers=num_workers,
                              shuffle=False,
                              pin_memory=True)
    test_loader = DataLoader(dataset=valid_dataset,
                             batch_size= batch_size,
                             num_workers=num_workers,
                             shuffle=False,
                             pin_memory=True)

    return train_loader, valid_loader, test_loader


if __name__ == '__main__':
    '''
    print(os.listdir('./data/windowed'))
    ds = IHDataset(root_dir='./data/windowed/', stage='valid')
    sample, target = ds[2]
    arr = np.transpose(255 * (0.5 * sample + 0.5), (1,2,0))
    im = Image.fromarray(np.uint8(arr))
    im.save('example.jpg')
    print(target)
    print(len(ds))
    '''
    ds = IHTestDataset(root_dir='../patients_windowed/001_1/')
    sample, img_path = ds[0]
    print('Img path:', img_path)
    arr = sample
    im = Image.fromarray(np.uint8(arr))
    im.save('example.jpg')