from torchvision.datasets.folder import pil_loader, accimage_loader, default_loader
from torch import Tensor
from pathlib import Path
from enum import Enum
from collections import namedtuple
from torchvision import transforms as T

import os
import numpy as np
import pdb
import functools

import torch.utils.data as data
import torch


class ConversionType(Enum):
    centerToVert = 1

def convert_bbox_format(boxes: Tensor, conversionType: int) -> Tensor:
    if conversionType > ConversionType.centerToVert.value:
        raise ValueError(
            f"conversionType must be less than" +
            "{ConversionType.centerToVert.value}, received {conversionType}")

    if conversionType == ConversionType.centerToVert.value:
        # convert box annotations from (Cx,Cy,W,H) to (X0,Y0,X1,Y1)
        box_centers = boxes[:, [0, 1, 0, 1]]
        box_wh = 0.5 * boxes[:, [2, 3, 2, 3]]
        box_wh[:, :2] *= -1
        convertedBoxes = box_centers + box_wh
    else:
        raise ValueError

    return convertedBoxes


class Wgisd(data.Dataset):
    """`FGVC-Aircraft <http://www.robots.ox.ac.uk/~vgg/data/fgvc-aircraft>`_ Dataset.
    Args:
        root (string): Root directory path to dataset.
        transform (callable, optional): A function/transform that takes in a PIL image
            and returns a transformed version. E.g. ``transforms.RandomCrop``
        loader (callable, optional): A function to load an image given its path.
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in the root directory. If dataset is already downloaded, it is not
            downloaded again.
    """
    url = 'https://github.com/thsant/wgisd.git'
    splits = ('train', 'test')

    def __init__(self, root, split='train', transform=None,
                 loader=default_loader, download=False,
                 val_size=0.2):
        if split not in self.splits:
            raise ValueError(
                'Split "{}" not found. Valid splits are: {}'.format(
                    split, ', '.join(
                        self.splits), ))
        if val_size < 0 or val_size > 1:
            raise ValueError('val_size should be a fraction between 0 and 1')
        self.root = Path(root)
        self.split = split

        # There's no file specifying a validation dataset, so use a subset of the
        # training dataset
        dset_file = self.split
        self.classes_file = self.root / f'{dset_file}.txt'

        if download:
            self.download()

        self.transform = transform
        self.loader = loader
        self.id_to_fname = {}
        self.val_size = val_size

        self.total_set = None
        self.samples = None
        self.create_dataset()
        self.mode = 'test' if self.split == 'test' else 'train'

    @property
    def mode(self):
        return self._mode

    @mode.setter
    def mode(self, mode):
        if self.split == 'test':
            self._mode = 'test'
            self.partition_dset()
            return

        supported_modes = ['train', 'val', 'trainval']
        if mode not in supported_modes:
            raise ValueError(f'mode must be one of {supported_modes}')

        self._mode = mode
        self.partition_dset()

    def create_dataset(self):
        image_names = []
        samples = []
        with open(self.classes_file, 'r') as f:
            for line in f:
                image_names.append(line.rstrip())
        data_dir = self.root / 'data'

        # Read bbox annotations from file
        for idx, img_name in enumerate(image_names):
            target = {}
            gt_boxes = []
            annotations = data_dir / f'{img_name}.txt'
            img_path = data_dir / f'{img_name}.jpg'

            with annotations.open() as f:
                for line in f:
                    gt_boxes.append(line.split()[1:])

            gt_np = np.array(gt_boxes, dtype=np.float32)
            gt_tensor = torch.as_tensor(gt_np, dtype=torch.float32)

            boxes = convert_bbox_format(gt_tensor, conversionType=1)
            img = self.loader(img_path)
            
            width, height = img.size
            boxes[:, [0, 2]] = boxes[:, [0, 2]] * width
            boxes[:, [1, 3]] = boxes[:, [1, 3]] * height
            boxes = boxes.to(dtype=torch.int32)
            numObjs = boxes.shape[0]
            labels = torch.ones((numObjs,), dtype=torch.int64)
            iscrowd = torch.zeros((numObjs,), dtype=torch.int64)
            image_id = torch.tensor([idx])
            self.id_to_fname[image_id.item()] = img_path.parts[-1]
            area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

            target['boxes'] = boxes
            target['labels'] = labels
            target['image_id'] = image_id
            target['area'] = area
            target['iscrowd'] = iscrowd

            samples.append((img_path, target))
        self.total_set = samples

    def partition_dset(self):
        num_images = len(self.total_set)
        split = int(np.floor(self.val_size * num_images))
        if self.mode == 'trainval':
            self.samples = self.total_set
        elif self.mode == 'train':
            self.samples = self.total_set[split:]
        elif self.mode == 'val':
            self.samples = self.total_set[:split]
        else:
            self.samples = self.total_set

    @functools.cached_property
    def mean(self):
        n_pixels = 0
        pix_sum = torch.zeros([3])
        for img_path, _ in self.total_set:
            img = self.loader(img_path)
            w,h = img.size
            im_tensor = T.ToTensor()(img)
            pix_sum += im_tensor.sum([1,2])
            n_pixels += (w*h)
        pix_avg = pix_sum / n_pixels
        return pix_avg

    @functools.cached_property
    def stddev(self):
        avg = self.mean
        avg = avg.reshape([3, 1, 1])
        var_sum = torch.zeros([3])
        n_pixels = 0
        for img_path, _ in self.total_set:
            img = self.loader(img_path)
            w,h = img.size
            im_tensor = T.ToTensor()(img)
            var_sum += ((im_tensor - avg)**2).sum([1,2])
            n_pixels += (w*h)
        var = var_sum / n_pixels
        return torch.sqrt(var)

    def get_fname(self, img_id):
        return self.id_to_fname[img_id.item()]

    def __getitem__(self, index):
        """
        Args:
            index (int): Index
        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """

        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample, target = self.transform(sample, target)

        return sample, target

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

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Root Location: {}\n'.format(self.root)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp,
                                     self.transform.__repr__().replace('\n',
                                                                       '\n' + ' ' * len(tmp)))
        return fmt_str

    def _check_exists(self):
        return self.root.exists() and self.classes_file.exists()

    def download(self):
        """Download the wgisd data if it doesn't exist already."""
        import requests
        import tarfile
        from git import Repo

        if self._check_exists():
            return

        print('Downloading %s ... (may take a few minutes)' % self.url)
        self.root.mkdir()
        Repo.clone_from(self.url, str(self.root))

        print('Done!')

[(62, 'pathlib.Path', 'Path', False, 'from pathlib import Path\n'), (161, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (175, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (184, 'torch.sqrt', 'torch.sqrt', False, 'import torch\n'), (120, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (121, 'torch.as_tensor', 'torch.as_tensor', False, 'import torch\n'), (131, 'torch.ones', 'torch.ones', False, 'import torch\n'), (132, 'torch.zeros', 'torch.zeros', False, 'import torch\n'), (133, 'torch.tensor', 'torch.tensor', False, 'import torch\n'), (148, 'numpy.floor', 'np.floor', True, 'import numpy as np\n'), (165, 'torchvision.transforms.ToTensor', 'T.ToTensor', True, 'from torchvision import transforms as T\n'), (180, 'torchvision.transforms.ToTensor', 'T.ToTensor', True, 'from torchvision import transforms as T\n')]