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)


[((233, 4, 233, 25), 'torch.manual_seed', 'torch.manual_seed', ({(233, 22, 233, 24): '(13)'}, {}), '(13)', False, 'import torch\n'), ((234, 18, 234, 56), 'torch.utils.data.random_split', 'random_split', ({(234, 31, 234, 33): 'ds', (234, 35, 234, 55): '[46, 46, 46, 45, 45]'}, {}), '(ds, [46, 46, 46, 45, 45])', False, 'from torch.utils.data import Subset, random_split, ConcatDataset\n'), ((18, 19, 18, 68), 'os.path.join', 'os.path.join', ({(18, 32, 18, 41): 'self.root', (18, 43, 18, 51): '"""images"""', (18, 53, 18, 67): 'self.imgs[idx]'}, {}), "(self.root, 'images', self.imgs[idx])", False, 'import os\n'), ((26, 16, 26, 59), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((29, 17, 29, 59), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((31, 19, 31, 38), 'torch.tensor', 'torch.tensor', ({(31, 32, 31, 37): '[idx]'}, {}), '([idx])', False, 'import torch\n'), ((68, 19, 68, 58), 'os.path.join', 'os.path.join', ({(68, 32, 68, 41): 'self.path', (68, 43, 68, 57): 'self.imgs[idx]'}, {}), '(self.path, self.imgs[idx])', False, 'import os\n'), ((76, 16, 76, 59), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((79, 17, 79, 59), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((81, 19, 81, 38), 'torch.tensor', 'torch.tensor', ({(81, 32, 81, 37): '[idx]'}, {}), '([idx])', False, 'import torch\n'), ((107, 20, 107, 45), 'os.path.join', 'os.path.join', ({(107, 33, 107, 37): 'root', (107, 39, 107, 44): '"""all"""'}, {}), "(root, 'all')", False, 'import os\n'), ((112, 19, 112, 58), 'os.path.join', 'os.path.join', ({(112, 32, 112, 41): 'self.path', (112, 43, 112, 57): 'self.imgs[idx]'}, {}), '(self.path, self.imgs[idx])', False, 'import os\n'), ((120, 16, 120, 59), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((123, 17, 123, 59), 'torch.as_tensor', 'torch.as_tensor', (), '', False, 'import torch\n'), ((125, 19, 125, 38), 'torch.tensor', 'torch.tensor', ({(125, 32, 125, 37): '[idx]'}, {}), '([idx])', False, 'import torch\n'), ((155, 19, 155, 68), 'os.path.join', 'os.path.join', ({(155, 32, 155, 41): 'self.root', (155, 43, 155, 51): '"""images"""', (155, 53, 155, 67): 'self.imgs[idx]'}, {}), "(self.root, 'images', self.imgs[idx])", False, 'import os\n'), ((157, 23, 157, 62), 'os.path.join', 'os.path.join', ({(157, 36, 157, 45): 'self.root', (157, 47, 157, 61): '"""density_maps"""'}, {}), "(self.root, 'density_maps')", False, 'import os\n'), ((159, 22, 159, 51), 'torch.from_numpy', 'torch.from_numpy', ({(159, 39, 159, 50): 'density_map'}, {}), '(density_map)', False, 'import torch\n'), ((15, 42, 15, 84), 'os.path.join', 'os.path.join', ({(15, 55, 15, 64): 'self.root', (15, 66, 15, 83): '"""annotations.csv"""'}, {}), "(self.root, 'annotations.csv')", False, 'import os\n'), ((58, 24, 58, 51), 'os.path.join', 'os.path.join', ({(58, 37, 58, 41): 'root', (58, 43, 58, 50): '"""train"""'}, {}), "(root, 'train')", False, 'import os\n'), ((60, 24, 60, 50), 'os.path.join', 'os.path.join', ({(60, 37, 60, 41): 'root', (60, 43, 60, 49): '"""test"""'}, {}), "(root, 'test')", False, 'import os\n'), ((64, 42, 64, 88), 'os.path.join', 'os.path.join', ({(64, 55, 64, 64): 'self.root', (64, 66, 64, 87): '"""annotations_all.csv"""'}, {}), "(self.root, 'annotations_all.csv')", False, 'import os\n'), ((109, 42, 109, 84), 'os.path.join', 'os.path.join', ({(109, 55, 109, 64): 'self.root', (109, 66, 109, 83): '"""annotations.csv"""'}, {}), "(self.root, 'annotations.csv')", False, 'import os\n'), ((158, 30, 158, 86), 'os.path.join', 'os.path.join', ({(158, 43, 158, 55): 'density_path', (158, 57, 158, 85): "self.imgs[idx][:-4] + '.npy'"}, {}), "(density_path, self.imgs[idx][:-4] + '.npy')", False, 'import os\n'), ((19, 14, 19, 34), 'PIL.Image.open', 'Image.open', ({(19, 25, 19, 33): 'img_path'}, {}), '(img_path)', False, 'from PIL import Image\n'), ((62, 32, 62, 53), 'os.listdir', 'os.listdir', ({(62, 43, 62, 52): 'self.path'}, {}), '(self.path)', False, 'import os\n'), ((69, 14, 69, 34), 'PIL.Image.open', 'Image.open', ({(69, 25, 69, 33): 'img_path'}, {}), '(img_path)', False, 'from PIL import Image\n'), ((108, 32, 108, 53), 'os.listdir', 'os.listdir', ({(108, 43, 108, 52): 'self.path'}, {}), '(self.path)', False, 'import os\n'), ((113, 14, 113, 34), 'PIL.Image.open', 'Image.open', ({(113, 25, 113, 33): 'img_path'}, {}), '(img_path)', False, 'from PIL import Image\n'), ((156, 14, 156, 34), 'PIL.Image.open', 'Image.open', ({(156, 25, 156, 33): 'img_path'}, {}), '(img_path)', False, 'from PIL import Image\n'), ((14, 43, 14, 71), 'os.path.join', 'os.path.join', ({(14, 56, 14, 60): 'root', (14, 62, 14, 70): '"""images"""'}, {}), "(root, 'images')", False, 'import os\n'), ((152, 43, 152, 71), 'os.path.join', 'os.path.join', ({(152, 56, 152, 60): 'root', (152, 62, 152, 70): '"""images"""'}, {}), "(root, 'images')", False, 'import os\n'), ((177, 27, 177, 56), 'os.path.join', 'os.path.join', ({(177, 40, 177, 44): 'root', (177, 46, 177, 55): '"""ans.npy"""'}, {}), "(root, 'ans.npy')", False, 'import os\n'), ((176, 44, 176, 60), 'os.listdir', 'os.listdir', ({(176, 55, 176, 59): 'root'}, {}), '(root)', False, 'import os\n')]