from __future__ import absolute_import, division, print_function

import re
from os import getcwd
from os.path import join

import numpy as np
import pandas as pd
import torch.utils.data as data
from PIL import Image

from torchvision.datasets.utils import check_integrity, download_url


class MuraDataset(data.Dataset):
    """`MURA <https://stanfordmlgroup.github.io/projects/mura/>`_ Dataset :
    Towards Radiologist-Level Abnormality Detection in Musculoskeletal Radiographs.
    """
    url = "https://cs.stanford.edu/group/mlgroup/mura-v1.0.zip"
    filename = "mura-v1.0.zip"
    md5_checksum = '4c36feddb7f5698c8bf291b912c438b1'
    _patient_re = re.compile(r'patient(\d+)')
    _study_re = re.compile(r'study(\d+)')
    _image_re = re.compile(r'image(\d+)')
    _study_type_re = re.compile(r'XR_(\w+)')

    def __init__(self, csv_f, transform=None, download=False):
        self.df = pd.read_csv(csv_f, names=['img', 'label'], header=None)
        self.imgs = self.df.img.values.tolist()
        self.labels = self.df.label.values.tolist()
        # following datasets/folder.py's weird convention here...
        self.samples = [tuple(x) for x in self.df.values]
        # number of unique classes
        self.classes = np.unique(self.labels)
        self.balanced_weights = self.balance_class_weights()

        self.transform = transform

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

    def _parse_patient(self, img_filename):
        return int(self._patient_re.search(img_filename).group(1))

    def _parse_study(self, img_filename):
        return int(self._study_re.search(img_filename).group(1))

    def _parse_image(self, img_filename):
        return int(self._image_re.search(img_filename).group(1))

    def _parse_study_type(self, img_filename):
        return self._study_type_re.search(img_filename).group(1)

    def download_and_uncompress_tarball(tarball_url, dataset_dir):
        """Downloads the `tarball_url` and uncompresses it locally.
        Args:
            tarball_url: The URL of a tarball file.
            dataset_dir: The directory where the temporary files are stored.
        """
        filename = tarball_url.split('/')[-1]
        filepath = os.path.join(dataset_dir, filename)

        def _progress(count, block_size, total_size):
            sys.stdout.write('\r>> Downloading %s %.1f%%' % (filename,
                                                             float(count * block_size) / float(total_size) * 100.0))
            sys.stdout.flush()

        filepath, _ = urllib.request.urlretrieve(tarball_url, filepath, _progress)
        print()
        statinfo = os.stat(filepath)
        print('Successfully downloaded', filename, statinfo.st_size, 'bytes.')
        if ".zip" in filename:
            print("zipfile:{}".format(filepath))
            with zipfile.ZipFile(filepath, "r") as zip_ref:
                zip_ref.extractall(dataset_dir)
        else:
            tarfile.open(filepath, 'r:gz').extractall(dataset_dir)

    def balance_class_weights(self):
        count = [0] * len(self.classes)
        for item in self.samples:
            count[item[1]] += 1
        weight_per_class = [0.] * len(self.classes)
        N = float(sum(count))
        for i in range(len(self.classes)):
            weight_per_class[i] = N / float(count[i])
        weight = [0] * len(self.samples)
        for idx, val in enumerate(self.samples):
            weight[idx] = weight_per_class[val[1]]
        return weight

    def __getitem__(self, idx):
        img_filename = join(self.imgs[idx])
        patient = self._parse_patient(img_filename)
        study = self._parse_study(img_filename)
        image_num = self._parse_image(img_filename)
        study_type = self._parse_study_type(img_filename)

        # todo(bdd) : inconsistent right now, need param for grayscale / RGB
        # todo(bdd) : 'L' -> gray, 'RGB' -> Colors
        image = Image.open(img_filename).convert('RGB')
        label = self.labels[idx]

        if self.transform is not None:
            image = self.transform(image)

        meta_data = {
            'y_true': label,
            'img_filename': img_filename,
            'patient': patient,
            'study': study,
            'study_type': study_type,
            'image_num': image_num,
            'encounter': "{}_{}_{}".format(study_type, patient, study)
        }
        return image, label, meta_data


if __name__ == '__main__':
    import torchvision.transforms as transforms
    import pprint

    data_dir = join(getcwd(), 'MURA-v1.0')
    val_csv = join(data_dir, 'valid.csv')
    val_loader = data.DataLoader(
        MuraDataset(val_csv,
                    transforms.Compose([
                        transforms.Resize(224),
                        transforms.CenterCrop(224),
                        transforms.ToTensor(),
                    ])),
        batch_size=1,
        shuffle=False,
        num_workers=1,
        pin_memory=False)

    for i, (image, label, meta_data) in enumerate(val_loader):
        pprint.pprint(meta_data.cpu())
        if i == 40:
            break