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, 'torch.manual_seed', 'torch.manual_seed', (['(13)'], {}), False, 'import torch\n'), (234, 'torch.utils.data.random_split', 'random_split', (['ds', '[46, 46, 46, 45, 45]'], {}), False, 'from torch.utils.data import Subset, random_split, ConcatDataset\n'), (18, 'os.path.join', 'os.path.join', (['self.root', '"""images"""', 'self.imgs[idx]'], {}), False, 'import os\n'), (26, 'torch.as_tensor', 'torch.as_tensor', (['boxes'], {'dtype': 'torch.float32'}), False, 'import torch\n'), (29, 'torch.as_tensor', 'torch.as_tensor', (['labels'], {'dtype': 'torch.int64'}), False, 'import torch\n'), (31, 'torch.tensor', 'torch.tensor', (['[idx]'], {}), False, 'import torch\n'), (68, 'os.path.join', 'os.path.join', (['self.path', 'self.imgs[idx]'], {}), False, 'import os\n'), (76, 'torch.as_tensor', 'torch.as_tensor', (['boxes'], {'dtype': 'torch.float32'}), False, 'import torch\n'), (79, 'torch.as_tensor', 'torch.as_tensor', (['labels'], {'dtype': 'torch.int64'}), False, 'import torch\n'), (81, 'torch.tensor', 'torch.tensor', (['[idx]'], {}), False, 'import torch\n'), (107, 'os.path.join', 'os.path.join', (['root', '"""all"""'], {}), False, 'import os\n'), (112, 'os.path.join', 'os.path.join', (['self.path', 'self.imgs[idx]'], {}), False, 'import os\n'), (120, 'torch.as_tensor', 'torch.as_tensor', (['boxes'], {'dtype': 'torch.float32'}), False, 'import torch\n'), (123, 'torch.as_tensor', 'torch.as_tensor', (['labels'], {'dtype': 'torch.int64'}), False, 'import torch\n'), (125, 'torch.tensor', 'torch.tensor', (['[idx]'], {}), False, 'import torch\n'), (155, 'os.path.join', 'os.path.join', (['self.root', '"""images"""', 'self.imgs[idx]'], {}), False, 'import os\n'), (157, 'os.path.join', 'os.path.join', (['self.root', '"""density_maps"""'], {}), False, 'import os\n'), (159, 'torch.from_numpy', 'torch.from_numpy', (['density_map'], {}), False, 'import torch\n'), (15, 'os.path.join', 'os.path.join', (['self.root', '"""annotations.csv"""'], {}), False, 'import os\n'), (58, 'os.path.join', 'os.path.join', (['root', '"""train"""'], {}), False, 'import os\n'), (60, 'os.path.join', 'os.path.join', (['root', '"""test"""'], {}), False, 'import os\n'), (64, 'os.path.join', 'os.path.join', (['self.root', '"""annotations_all.csv"""'], {}), False, 'import os\n'), (109, 'os.path.join', 'os.path.join', (['self.root', '"""annotations.csv"""'], {}), False, 'import os\n'), (158, 'os.path.join', 'os.path.join', (['density_path', "(self.imgs[idx][:-4] + '.npy')"], {}), False, 'import os\n'), (19, 'PIL.Image.open', 'Image.open', (['img_path'], {}), False, 'from PIL import Image\n'), (62, 'os.listdir', 'os.listdir', (['self.path'], {}), False, 'import os\n'), (69, 'PIL.Image.open', 'Image.open', (['img_path'], {}), False, 'from PIL import Image\n'), (108, 'os.listdir', 'os.listdir', (['self.path'], {}), False, 'import os\n'), (113, 'PIL.Image.open', 'Image.open', (['img_path'], {}), False, 'from PIL import Image\n'), (156, 'PIL.Image.open', 'Image.open', (['img_path'], {}), False, 'from PIL import Image\n'), (14, 'os.path.join', 'os.path.join', (['root', '"""images"""'], {}), False, 'import os\n'), (152, 'os.path.join', 'os.path.join', (['root', '"""images"""'], {}), False, 'import os\n'), (177, 'os.path.join', 'os.path.join', (['root', '"""ans.npy"""'], {}), False, 'import os\n'), (176, 'os.listdir', 'os.listdir', (['root'], {}), False, 'import os\n')]