import numpy as np
import os
import pandas as pd
import torch
from PIL import Image, ImageFile
import variables
from utils import iou_width_height as iou, NMS, plot_image, CTF
from torch.utils.data import Dataset, DataLoader

ImageFile.LOAD_TRUNCATED_IMAGES = True


class KittiDataset(torch.utils.data.Dataset):
    OBJECTS = ['DontCare', 'Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc']

    def __init__(self, filename, image_dir, label_dir, anchors, split_size=[13, 26, 52], no_classes=9, transform=None):
        self.connections = pd.read_csv(filename)
        self.label_dir = label_dir
        self.image_dir = image_dir
        self.split_size = split_size
        self.transform = transform
        self.anchors = torch.tensor(anchors[0] + anchors[1] + anchors[2])
        self.no_anchors = self.anchors.shape[0]
        self.no_anchors_scale = self.no_anchors // 3
        self.no_classes = no_classes
        self.ignore_iou_lim = 0.5

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

    def __getitem__(self, item):
        labels_path = os.path.join(self.label_dir, self.connections.iloc[item, 0])
        image_path = os.path.join(self.image_dir, self.connections.iloc[item, 1])
        img = Image.open(image_path).convert("RGB")
        img_arr = np.array(img)
        pic_width, pic_height = img.size
        frames = []

        with open(labels_path) as l:
            for label in l.readlines():
                try:
                    object_type = self.OBJECTS[[label.find(element) for element in self.OBJECTS].index(0)]
                    truncated_label = label[[i for i, x in enumerate(label) if x == " "][3]:
                                            [i for i, x in enumerate(label) if x == " "][7]]

                    class_label, left, top, right, bottom = (object_type + truncated_label).split(' ')
                    # normalize picture width base on picture size
                    x, y, width, height = ((float(left) + float(right)) / 2) / pic_width, \
                                          ((float(top) + float(bottom)) / 2) / pic_height, \
                                          (float(right) - float(left)) / pic_width, \
                                          (float(bottom) - float(top)) / pic_height

                    frames.append([x, y, width, height, int(self.OBJECTS.index(class_label))])
                except:
                    pass
        if self.transform:
            try:
                transforms = self.transform(image=img_arr, bboxes=frames)
                img_arr = transforms["image"]
                frames = transforms["bboxes"]
            except:
                print('\n Image or file corrupted')
                img_arr = img_arr
                frames = frames
                # return (torch.randn((3, variables.IMG_SHAPE, variables.IMG_SHAPE)),
                #        tuple([torch.zeros((self.no_anchors // 3, i, i, 6)) for i in self.split_size]))

        destinations = [torch.zeros((self.no_anchors // 3, i, i, 6)) for i in self.split_size]
        for frame in frames:
            anchors = iou(torch.tensor(frame[2:4]), self.anchors)
            anchor_dets = anchors.argsort(descending=True, dim=0)
            x, y, width, height, class_label = frame
            has_anchor = [False, False, False]
            for index in anchor_dets:
                scale_idx = index // self.no_anchors_scale
                anchor_on_scale = index % self.no_anchors_scale
                split = self.split_size[scale_idx]
                a, b = int(split * y), int(split * x)  # which cell
                anchor_available = destinations[scale_idx][anchor_on_scale, a, b, 0]
                if not anchor_available and not has_anchor[scale_idx]:
                    destinations[scale_idx][anchor_on_scale, a, b, 0] = 1
                    x_cell, y_cell = split * x - b, split * y - a  # both between [0,1]
                    width_cell, height_cell = (
                        width * split,
                        height * split,
                    )  # can be greater than 1 since it's relative to cell
                    box_coordinates = torch.tensor(
                        [x_cell, y_cell, width_cell, height_cell]
                    )
                    destinations[scale_idx][anchor_on_scale, a, b, 1:5] = box_coordinates
                    destinations[scale_idx][anchor_on_scale, a, b, 5] = int(class_label)
                    has_anchor[scale_idx] = True

                elif not anchor_available and anchors[index] > self.ignore_iou_lim:
                    destinations[scale_idx][anchor_on_scale, a, b, 0] = -1  # ignore prediction
        return img_arr, tuple(destinations)

# TODO: Delete before ending
def test():
    anchors = variables.ANCHORS

    transform = variables.test_transforms

    dataset = KittiDataset(
        "data/kitti_single/training_data.csv",
        "data/kitti_single/training/image_2",
        "data/kitti_single/training/label_2",
        split_size=[13, 26, 52],
        anchors=anchors,
        transform=transform,
    )
    S = [13, 26, 52]
    scaled_anchors = torch.tensor(anchors) / (
            1 / torch.tensor(S).unsqueeze(1).unsqueeze(1).repeat(1, 3, 2)
    )
    loader = DataLoader(dataset=dataset, batch_size=1, shuffle=True)
    for x, y in loader:
        boxes = []

        for i in range(y[0].shape[1]):
            anchor = scaled_anchors[i]
            # print(anchor.shape)
            # print(y[i].shape)
            boxes += CTF(
                y[i], is_preds=False, SPLIT_SIZE=y[i].shape[2], anchors=anchor
            )[0]
        boxes = NMS(boxes, iou_lim=1, prob_lim=0.7, frame_format="1_point")
        # print(boxes)
        plot_image(x[0].permute(1, 2, 0).to("cpu"), boxes)


if __name__ == '__main__':
    test()
