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, 20, 62, 30), 'pathlib.Path', 'Path', ({(62, 25, 62, 29): 'root'}, {}), '(root)', False, 'from pathlib import Path\n'), ((161, 18, 161, 34), 'torch.zeros', 'torch.zeros', ({(161, 30, 161, 33): '[3]'}, {}), '([3])', False, 'import torch\n'), ((175, 18, 175, 34), 'torch.zeros', 'torch.zeros', ({(175, 30, 175, 33): '[3]'}, {}), '([3])', False, 'import torch\n'), ((184, 15, 184, 30), 'torch.sqrt', 'torch.sqrt', ({(184, 26, 184, 29): 'var'}, {}), '(var)', False, 'import torch\n'), ((120, 20, 120, 56), 'numpy.array', 'np.array', (), '', True, 'import numpy as np\n'), ((121, 24, 121, 67), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((131, 21, 131, 62), 'torch.ones', 'torch.ones', (), '', False, 'import torch\n'), ((132, 22, 132, 64), 'torch.zeros', 'torch.zeros', (), '', False, 'import torch\n'), ((133, 23, 133, 42), 'torch.tensor', 'torch.tensor', ({(133, 36, 133, 41): '[idx]'}, {}), '([idx])', False, 'import torch\n'), ((148, 20, 148, 56), 'numpy.floor', 'np.floor', ({(148, 29, 148, 55): 'self.val_size * num_images'}, {}), '(self.val_size * num_images)', True, 'import numpy as np\n'), ((165, 24, 165, 36), 'torchvision.transforms.ToTensor', 'T.ToTensor', ({}, {}), '()', True, 'from torchvision import transforms as T\n'), ((180, 24, 180, 36), 'torchvision.transforms.ToTensor', 'T.ToTensor', ({}, {}), '()', True, 'from torchvision import transforms as T\n')]