import os
import tarfile

import h5py
import numpy as np
import torch
import torch.utils.data as data
from PIL import Image
from scipy import io
from torchvision.datasets.utils import download_url

DATASET_YEAR_DICT = {
    "2012": {
        "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar",
        "filename": "VOCtrainval_11-May-2012.tar",
        "md5": "6cd6e144f989b92b3379bac3b3de84fd",
        "base_dir": "VOCdevkit/VOC2012",
    },
    "2011": {
        "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2011/VOCtrainval_25-May-2011.tar",
        "filename": "VOCtrainval_25-May-2011.tar",
        "md5": "6c3384ef61512963050cb5d687e5bf1e",
        "base_dir": "TrainVal/VOCdevkit/VOC2011",
    },
    "2010": {
        "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2010/VOCtrainval_03-May-2010.tar",
        "filename": "VOCtrainval_03-May-2010.tar",
        "md5": "da459979d0c395079b5c75ee67908abb",
        "base_dir": "VOCdevkit/VOC2010",
    },
    "2009": {
        "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2009/VOCtrainval_11-May-2009.tar",
        "filename": "VOCtrainval_11-May-2009.tar",
        "md5": "59065e4b188729180974ef6572f6a212",
        "base_dir": "VOCdevkit/VOC2009",
    },
    "2008": {
        "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2008/VOCtrainval_14-Jul-2008.tar",
        "filename": "VOCtrainval_11-May-2012.tar",
        "md5": "2629fa636546599198acfcfbfcf1904a",
        "base_dir": "VOCdevkit/VOC2008",
    },
    "2007": {
        "url": "http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar",
        "filename": "VOCtrainval_06-Nov-2007.tar",
        "md5": "c52e279531787c972589f7e41ab4ae64",
        "base_dir": "VOCdevkit/VOC2007",
    },
}


class VOCSegmentation(data.Dataset):
    """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Segmentation Dataset.

    Args:
        root (string): Root directory of the VOC Dataset.
        year (string, optional): The dataset year, supports years 2007 to 2012.
        image_set (string, optional): Select the image_set to use, ``train``, ``trainval`` or ``val``
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
    """

    CLASSES = 20
    CLASSES_NAMES = [
        "aeroplane",
        "bicycle",
        "bird",
        "boat",
        "bottle",
        "bus",
        "car",
        "cat",
        "chair",
        "cow",
        "diningtable",
        "dog",
        "horse",
        "motorbike",
        "person",
        "potted-plant",
        "sheep",
        "sofa",
        "train",
        "tvmonitor",
        "ambigious",
    ]

    def __init__(
        self,
        root,
        year="2012",
        image_set="train",
        download=False,
        transform=None,
        target_transform=None,
    ):
        self.root = os.path.expanduser(root)
        self.year = year
        self.url = DATASET_YEAR_DICT[year]["url"]
        self.filename = DATASET_YEAR_DICT[year]["filename"]
        self.md5 = DATASET_YEAR_DICT[year]["md5"]
        self.transform = transform
        self.target_transform = target_transform
        self.image_set = image_set
        base_dir = DATASET_YEAR_DICT[year]["base_dir"]
        voc_root = os.path.join(self.root, base_dir)
        image_dir = os.path.join(voc_root, "JPEGImages")
        mask_dir = os.path.join(voc_root, "SegmentationClass")

        if download:
            download_extract(self.url, self.root, self.filename, self.md5)

        if not os.path.isdir(voc_root):
            raise RuntimeError(
                "Dataset not found or corrupted."
                + " You can use download=True to download it"
            )

        splits_dir = os.path.join(voc_root, "ImageSets/Segmentation")

        split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt")

        if not os.path.exists(split_f):
            raise ValueError(
                'Wrong image_set entered! Please use image_set="train" '
                'or image_set="trainval" or image_set="val"'
            )

        with open(os.path.join(split_f), "r") as f:
            file_names = [x.strip() for x in f.readlines()]

        self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names]
        self.masks = [os.path.join(mask_dir, x + ".png") for x in file_names]
        assert len(self.images) == len(self.masks)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is the image segmentation.
        """
        img = Image.open(self.images[index]).convert("RGB")
        target = Image.open(self.masks[index])

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

        if self.target_transform is not None:
            target = np.array(self.target_transform(target)).astype("int32")
            target[target == 255] = -1
            target = torch.from_numpy(target).long()

        return img, target

    @staticmethod
    def _mask_transform(mask):
        target = np.array(mask).astype("int32")
        target[target == 255] = -1
        return torch.from_numpy(target).long()

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

    @property
    def pred_offset(self):
        return 0


class VOCClassification(data.Dataset):
    """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Segmentation Dataset.

    Args:
        root (string): Root directory of the VOC Dataset.
        year (string, optional): The dataset year, supports years 2007 to 2012.
        image_set (string, optional): Select the image_set to use, ``train``, ``trainval`` or ``val``
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
    """

    CLASSES = 20

    def __init__(
        self, root, year="2012", image_set="train", download=False, transform=None
    ):
        self.root = os.path.expanduser(root)
        self.year = year
        self.url = DATASET_YEAR_DICT[year]["url"]
        self.filename = DATASET_YEAR_DICT[year]["filename"]
        self.md5 = DATASET_YEAR_DICT[year]["md5"]
        self.transform = transform
        self.image_set = image_set
        base_dir = DATASET_YEAR_DICT[year]["base_dir"]
        voc_root = os.path.join(self.root, base_dir)
        image_dir = os.path.join(voc_root, "JPEGImages")
        mask_dir = os.path.join(voc_root, "SegmentationClass")

        if download:
            download_extract(self.url, self.root, self.filename, self.md5)

        if not os.path.isdir(voc_root):
            raise RuntimeError(
                "Dataset not found or corrupted."
                + " You can use download=True to download it"
            )

        splits_dir = os.path.join(voc_root, "ImageSets/Segmentation")

        split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt")

        if not os.path.exists(split_f):
            raise ValueError(
                'Wrong image_set entered! Please use image_set="train" '
                'or image_set="trainval" or image_set="val"'
            )

        with open(os.path.join(split_f), "r") as f:
            file_names = [x.strip() for x in f.readlines()]

        self.images = [os.path.join(image_dir, x + ".jpg") for x in file_names]
        self.masks = [os.path.join(mask_dir, x + ".png") for x in file_names]
        assert len(self.images) == len(self.masks)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is the image segmentation.
        """
        img = Image.open(self.images[index]).convert("RGB")
        target = Image.open(self.masks[index])

        # if self.transform is not None:
        #     img = self.transform(img)
        if self.transform is not None:
            img, target = self.transform(img, target)

        visible_classes = np.unique(target)
        labels = torch.zeros(self.CLASSES)
        for id in visible_classes:
            if id not in (0, 255):
                labels[id - 1].fill_(1)

        return img, labels

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


class VOCSBDClassification(data.Dataset):
    """`Pascal VOC <http://host.robots.ox.ac.uk/pascal/VOC/>`_ Segmentation Dataset.

    Args:
        root (string): Root directory of the VOC Dataset.
        year (string, optional): The dataset year, supports years 2007 to 2012.
        image_set (string, optional): Select the image_set to use, ``train``, ``trainval`` or ``val``
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in root directory. If dataset is already downloaded, it is not
            downloaded again.
        transform (callable, optional): A function/transform that  takes in an PIL image
            and returns a transformed version. E.g, ``transforms.RandomCrop``
    """

    CLASSES = 20

    def __init__(
        self,
        root,
        sbd_root,
        year="2012",
        image_set="train",
        download=False,
        transform=None,
    ):
        self.root = os.path.expanduser(root)
        self.sbd_root = os.path.expanduser(sbd_root)
        self.year = year
        self.url = DATASET_YEAR_DICT[year]["url"]
        self.filename = DATASET_YEAR_DICT[year]["filename"]
        self.md5 = DATASET_YEAR_DICT[year]["md5"]
        self.transform = transform
        self.image_set = image_set
        base_dir = DATASET_YEAR_DICT[year]["base_dir"]
        voc_root = os.path.join(self.root, base_dir)
        image_dir = os.path.join(voc_root, "JPEGImages")
        mask_dir = os.path.join(voc_root, "SegmentationClass")
        sbd_image_dir = os.path.join(sbd_root, "img")
        sbd_mask_dir = os.path.join(sbd_root, "cls")

        if download:
            download_extract(self.url, self.root, self.filename, self.md5)

        if not os.path.isdir(voc_root):
            raise RuntimeError(
                "Dataset not found or corrupted."
                + " You can use download=True to download it"
            )

        splits_dir = os.path.join(voc_root, "ImageSets/Segmentation")

        split_f = os.path.join(splits_dir, image_set.rstrip("\n") + ".txt")
        sbd_split = os.path.join(sbd_root, "train.txt")

        if not os.path.exists(split_f):
            raise ValueError(
                'Wrong image_set entered! Please use image_set="train" '
                'or image_set="trainval" or image_set="val"'
            )

        with open(os.path.join(split_f), "r") as f:
            voc_file_names = [x.strip() for x in f.readlines()]

        with open(os.path.join(sbd_split), "r") as f:
            sbd_file_names = [x.strip() for x in f.readlines()]

        self.images = [os.path.join(image_dir, x + ".jpg") for x in voc_file_names]
        self.images += [os.path.join(sbd_image_dir, x + ".jpg") for x in sbd_file_names]
        self.masks = [os.path.join(mask_dir, x + ".png") for x in voc_file_names]
        self.masks += [os.path.join(sbd_mask_dir, x + ".mat") for x in sbd_file_names]
        assert len(self.images) == len(self.masks)

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is the image segmentation.
        """
        img = Image.open(self.images[index]).convert("RGB")
        mask_path = self.masks[index]
        if mask_path[-3:] == "mat":
            target = io.loadmat(mask_path, struct_as_record=False, squeeze_me=True)[
                "GTcls"
            ].Segmentation
            target = Image.fromarray(target, mode="P")
        else:
            target = Image.open(self.masks[index])

        if self.transform is not None:
            img, target = self.transform(img, target)

        visible_classes = np.unique(target)
        labels = torch.zeros(self.CLASSES)
        for id in visible_classes:
            if id not in (0, 255):
                labels[id - 1].fill_(1)

        return img, labels

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


def download_extract(url, root, filename, md5):
    download_url(url, root, filename, md5)
    with tarfile.open(os.path.join(root, filename), "r") as tar:
        tar.extractall(path=root)


class VOCResults(data.Dataset):
    CLASSES = 20
    CLASSES_NAMES = [
        "aeroplane",
        "bicycle",
        "bird",
        "boat",
        "bottle",
        "bus",
        "car",
        "cat",
        "chair",
        "cow",
        "diningtable",
        "dog",
        "horse",
        "motorbike",
        "person",
        "potted-plant",
        "sheep",
        "sofa",
        "train",
        "tvmonitor",
        "ambigious",
    ]

    def __init__(self, path):
        super(VOCResults, self).__init__()

        self.path = os.path.join(path, "results.hdf5")
        self.data = None

        print("Reading dataset length...")
        with h5py.File(self.path, "r") as f:
            self.data_length = len(f["/image"])

    def __len__(self):
        return self.data_length

    def __getitem__(self, item):
        if self.data is None:
            self.data = h5py.File(self.path, "r")

        image = torch.tensor(self.data["image"][item])
        vis = torch.tensor(self.data["vis"][item])
        target = torch.tensor(self.data["target"][item])
        class_pred = torch.tensor(self.data["class_pred"][item])

        return image, vis, target, class_pred