import os
import torch
from PIL import Image
from read_csv import csv_to_label_and_bbx
import numpy as np
from torch.utils.data import Subset, random_split, ConcatDataset


class NBIDataset(object):
    def __init__(self, root, transforms, nob3=False):
        self.root = root
        self.transforms = transforms
        # load all image files, sorting them to ensure that they are aligned
        self.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))
        self.boxes = csv_to_label_and_bbx(os.path.join(self.root, "annotations.csv"), nob3)

    def __getitem__(self, idx):
        img_path = os.path.join(self.root, "images", self.imgs[idx])
        img = Image.open(img_path).convert("RGB")

        annotations = self.boxes[self.imgs[idx]]
        boxes = annotations['bbx']
        labels = annotations['labels']

        # FloatTensor[N, 4]
        boxes = torch.as_tensor(boxes, dtype=torch.float32)

        # Int64Tensor[N]
        labels = torch.as_tensor(labels, dtype=torch.int64)

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((labels.size()[0],), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        # target["image_path"] = img_path
        target["area"] = area
        target["iscrowd"] = iscrowd

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

        return img, target

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


class NBINewDataset(object):
    def __init__(self, root, transforms, train=True):
        self.root = root
        self.transforms = transforms
        if train:
            self.path = os.path.join(root, "train")
        else:
            self.path = os.path.join(root, "test")

        self.imgs = list(sorted(os.listdir(self.path)))

        self.boxes = csv_to_label_and_bbx(os.path.join(self.root, "annotations_all.csv"), img_names=self.imgs)


    def __getitem__(self, idx):
        img_path = os.path.join(self.path, self.imgs[idx])
        img = Image.open(img_path).convert("RGB")

        annotations = self.boxes[self.imgs[idx]]
        boxes = annotations['bbx']
        labels = annotations['labels']

        # FloatTensor[N, 4]
        boxes = torch.as_tensor(boxes, dtype=torch.float32)

        # Int64Tensor[N]
        labels = torch.as_tensor(labels, dtype=torch.int64)

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((labels.size()[0],), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        # target["image_path"] = img_path
        target["area"] = area
        target["iscrowd"] = iscrowd

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

        return img, target

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


class NBIFullDataset(object):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        self.path = os.path.join(root, "all")
        self.imgs = list(sorted(os.listdir(self.path)))
        self.boxes = csv_to_label_and_bbx(os.path.join(self.root, "annotations.csv"), img_names=self.imgs)

    def __getitem__(self, idx):
        img_path = os.path.join(self.path, self.imgs[idx])
        img = Image.open(img_path).convert("RGB")

        annotations = self.boxes[self.imgs[idx]]
        boxes = annotations['bbx']
        labels = annotations['labels']

        # FloatTensor[N, 4]
        boxes = torch.as_tensor(boxes, dtype=torch.float32)

        # Int64Tensor[N]
        labels = torch.as_tensor(labels, dtype=torch.int64)

        image_id = torch.tensor([idx])
        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])
        iscrowd = torch.zeros((labels.size()[0],), dtype=torch.int64)

        target = {}
        target["boxes"] = boxes
        target["labels"] = labels
        target["image_id"] = image_id
        # target["image_path"] = img_path
        target["area"] = area
        target["iscrowd"] = iscrowd

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

        return img, target

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


class NBIDenseDataset(object):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        # load all image files, sorting them to ensure that they are aligned
        self.imgs = list(sorted(os.listdir(os.path.join(root, "images"))))

    def __getitem__(self, idx):
        img_path = os.path.join(self.root, "images", self.imgs[idx])
        img = Image.open(img_path).convert("RGB")
        density_path = os.path.join(self.root, "density_maps")
        density_map = np.load(os.path.join(density_path, self.imgs[idx][:-4] + ".npy"))
        density_map = torch.from_numpy(density_map)

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

        return img, density_map

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


class NBIPatchDataset(object):
    def __init__(self, root, transforms):
        self.root = root
        self.transforms = transforms
        # load all image files, sorting them to ensure that they are aligned
        self.imgs = [x for x in list(sorted(os.listdir(root))) if x[-3:] == "png"]
        self.ans = np.load(os.path.join(root, "ans.npy"), allow_pickle=True).item()

    def __getitem__(self, idx):
        # img_path = os.path.join(self.root, "images", self.imgs[idx])
        # img = Image.open(img_path).convert("RGB")
        # density_path = os.path.join(self.root, "density_maps")
        # density_map = np.load(os.path.join(density_path, self.imgs[idx][:-4] + ".npy"))
        # density_map = torch.from_numpy(density_map)
        #
        # if self.transforms is not None:
        #     img = self.transforms(img)
        #     # target = self.transforms(target)

        return self.imgs[idx]

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


def split_index(K=5, len=100):
    idx = list(range(len))
    final_list = []
    for i in range(K):
        final_list.append(idx[(i*len)//K:((i+1)*len)//K])
    return final_list


def k_fold_index(K=5, len=100, fold=0):
    split = split_index(K, len)
    val = split[fold]
    train = []
    for i in range(K):
        if i != fold:
            train = train + split[i]
    return train, val


def stat_dataset(dataset):
    class_ids = {1: "A", 2: "B1", 3: "B2", 4: "B3"}
    stats = {"A": 0, "B1": 0, "B2": 0, "B3": 0}
    for img, target in dataset:
        for k in target['labels']:
            stats[class_ids[int(k)]] += 1
    print(stats)


def  NBIFiveFoldDataset(transforms):
    ds = NBIFullDataset(root="./NBI_full_dataset/", transforms=transforms)
    # n = len(ds)
    # for i in range(5):
    #     train_idx, val_idx = k_fold_index(5, n, i)
    #     train_subset = Subset(ds, train_idx)
    #     val_subset = Subset(ds, val_idx)
    #     print("Fold: %d" % i, len(train_subset), len(val_subset))
    #     stat_dataset(train_subset)
    #     stat_dataset(val_subset)
    torch.manual_seed(13)
    all_subsets = random_split(ds, [46, 46, 46, 45, 45])
    fold_i_subsets = []
    for i in range(5):
        val_subset = all_subsets[i]
        train_subset = ConcatDataset([all_subsets[j] for j in range(5) if j != i])
        fold_i_subsets.append({"train": train_subset, "val": val_subset})
        # print("Fold: %d" % i, len(train_subset), len(val_subset))
        # stat_dataset(train_subset)
        # stat_dataset(val_subset)
    return fold_i_subsets

if __name__ == '__main__':
    # ds = NBIFiveFoldDataset(None)
    di = "aaa".encode("UTF-8")
    result = eval(di)
    print(result)


[((7492, 7513), 'torch.manual_seed', 'torch.manual_seed', (['(13)'], {}), '(13)\n', (7509, 7513), False, 'import torch\n'), ((7532, 7570), 'torch.utils.data.random_split', 'random_split', (['ds', '[46, 46, 46, 45, 45]'], {}), '(ds, [46, 46, 46, 45, 45])\n', (7544, 7570), False, 'from torch.utils.data import Subset, random_split, ConcatDataset\n'), ((611, 660), 'os.path.join', 'os.path.join', (['self.root', '"""images"""', 'self.imgs[idx]'], {}), "(self.root, 'images', self.imgs[idx])\n", (623, 660), False, 'import os\n'), ((880, 923), 'torch.as_tensor', 'torch.as_tensor', (['boxes'], {'dtype': 'torch.float32'}), '(boxes, dtype=torch.float32)\n', (895, 923), False, 'import torch\n'), ((967, 1009), 'torch.as_tensor', 'torch.as_tensor', (['labels'], {'dtype': 'torch.int64'}), '(labels, dtype=torch.int64)\n', (982, 1009), False, 'import torch\n'), ((1030, 1049), 'torch.tensor', 'torch.tensor', (['[idx]'], {}), '([idx])\n', (1042, 1049), False, 'import torch\n'), ((2139, 2178), 'os.path.join', 'os.path.join', (['self.path', 'self.imgs[idx]'], {}), '(self.path, self.imgs[idx])\n', (2151, 2178), False, 'import os\n'), ((2398, 2441), 'torch.as_tensor', 'torch.as_tensor', (['boxes'], {'dtype': 'torch.float32'}), '(boxes, dtype=torch.float32)\n', (2413, 2441), False, 'import torch\n'), ((2485, 2527), 'torch.as_tensor', 'torch.as_tensor', (['labels'], {'dtype': 'torch.int64'}), '(labels, dtype=torch.int64)\n', (2500, 2527), False, 'import torch\n'), ((2548, 2567), 'torch.tensor', 'torch.tensor', (['[idx]'], {}), '([idx])\n', (2560, 2567), False, 'import torch\n'), ((3309, 3334), 'os.path.join', 'os.path.join', (['root', '"""all"""'], {}), "(root, 'all')\n", (3321, 3334), False, 'import os\n'), ((3550, 3589), 'os.path.join', 'os.path.join', (['self.path', 'self.imgs[idx]'], {}), '(self.path, self.imgs[idx])\n', (3562, 3589), False, 'import os\n'), ((3809, 3852), 'torch.as_tensor', 'torch.as_tensor', (['boxes'], {'dtype': 'torch.float32'}), '(boxes, dtype=torch.float32)\n', (3824, 3852), False, 'import torch\n'), ((3896, 3938), 'torch.as_tensor', 'torch.as_tensor', (['labels'], {'dtype': 'torch.int64'}), '(labels, dtype=torch.int64)\n', (3911, 3938), False, 'import torch\n'), ((3959, 3978), 'torch.tensor', 'torch.tensor', (['[idx]'], {}), '([idx])\n', (3971, 3978), False, 'import torch\n'), ((4905, 4954), 'os.path.join', 'os.path.join', (['self.root', '"""images"""', 'self.imgs[idx]'], {}), "(self.root, 'images', self.imgs[idx])\n", (4917, 4954), False, 'import os\n'), ((5028, 5067), 'os.path.join', 'os.path.join', (['self.root', '"""density_maps"""'], {}), "(self.root, 'density_maps')\n", (5040, 5067), False, 'import os\n'), ((5178, 5207), 'torch.from_numpy', 'torch.from_numpy', (['density_map'], {}), '(density_map)\n', (5194, 5207), False, 'import torch\n'), ((509, 551), 'os.path.join', 'os.path.join', (['self.root', '"""annotations.csv"""'], {}), "(self.root, 'annotations.csv')\n", (521, 551), False, 'import os\n'), ((1824, 1851), 'os.path.join', 'os.path.join', (['root', '"""train"""'], {}), "(root, 'train')\n", (1836, 1851), False, 'import os\n'), ((1890, 1916), 'os.path.join', 'os.path.join', (['root', '"""test"""'], {}), "(root, 'test')\n", (1902, 1916), False, 'import os\n'), ((2017, 2063), 'os.path.join', 'os.path.join', (['self.root', '"""annotations_all.csv"""'], {}), "(self.root, 'annotations_all.csv')\n", (2029, 2063), False, 'import os\n'), ((3433, 3475), 'os.path.join', 'os.path.join', (['self.root', '"""annotations.csv"""'], {}), "(self.root, 'annotations.csv')\n", (3445, 3475), False, 'import os\n'), ((5098, 5154), 'os.path.join', 'os.path.join', (['density_path', "(self.imgs[idx][:-4] + '.npy')"], {}), "(density_path, self.imgs[idx][:-4] + '.npy')\n", (5110, 5154), False, 'import os\n'), ((675, 695), 'PIL.Image.open', 'Image.open', (['img_path'], {}), '(img_path)\n', (685, 695), False, 'from PIL import Image\n'), ((1950, 1971), 'os.listdir', 'os.listdir', (['self.path'], {}), '(self.path)\n', (1960, 1971), False, 'import os\n'), ((2193, 2213), 'PIL.Image.open', 'Image.open', (['img_path'], {}), '(img_path)\n', (2203, 2213), False, 'from PIL import Image\n'), ((3367, 3388), 'os.listdir', 'os.listdir', (['self.path'], {}), '(self.path)\n', (3377, 3388), False, 'import os\n'), ((3604, 3624), 'PIL.Image.open', 'Image.open', (['img_path'], {}), '(img_path)\n', (3614, 3624), False, 'from PIL import Image\n'), ((4969, 4989), 'PIL.Image.open', 'Image.open', (['img_path'], {}), '(img_path)\n', (4979, 4989), False, 'from PIL import Image\n'), ((435, 463), 'os.path.join', 'os.path.join', (['root', '"""images"""'], {}), "(root, 'images')\n", (447, 463), False, 'import os\n'), ((4821, 4849), 'os.path.join', 'os.path.join', (['root', '"""images"""'], {}), "(root, 'images')\n", (4833, 4849), False, 'import os\n'), ((5746, 5775), 'os.path.join', 'os.path.join', (['root', '"""ans.npy"""'], {}), "(root, 'ans.npy')\n", (5758, 5775), False, 'import os\n'), ((5680, 5696), 'os.listdir', 'os.listdir', (['root'], {}), '(root)\n', (5690, 5696), False, 'import os\n')]