import torch

import torch.nn.functional as F

from utils import *

from PIL import Image, ImageDraw, ImageFont

from utils.load_ckpt import load_model

from utils.ssd_eval_utils import *

from model import build_ssd

from configs import SSD as opt





def detect_pill(original_image, preprocess_image, model, min_score, max_overlap, top_k, suppress=None, visualize=False):

    """

    Detect objects in an image with a trained SSD300, and visualize the results.

    :param original_image: image, a PIL Image

    :param min_score: minimum threshold for a detected box to be considered a match for a certain class

    :param max_overlap: maximum overlap two boxes can have so that the one with the lower score is not suppressed via Non-Maximum Suppression (NMS)

    :param top_k: if there are a lot of resulting detection across all classes, keep only the top 'k'

    :param suppress: classes that you know for sure cannot be in the image or you do not want in the image, a list

    :return: annotated image, a PIL Image

    """

    # define variables

    labelmap = opt['labelmap']

    rev_label_map = {v: k for k, v in labelmap.items()}

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    res_path = './result/'



    # Move to default device

    input_image = preprocess_image.to(device)



    # Forward prop.

    predicted_locs, predicted_scores, priors = model(input_image)



    # Detect objects in SSD output

    det_boxes, det_labels, det_scores = detect_objects(predicted_locs, predicted_scores, priors, min_score=min_score,

                                                             max_overlap=max_overlap, top_k=top_k)



    # Move detections to the CPU

    det_boxes = det_boxes[0].to('cpu')



    # Transform to original image dimensions

    original_dims = torch.FloatTensor(

        [original_image.width, original_image.height, original_image.width, original_image.height]).unsqueeze(0)

    det_boxes = det_boxes * original_dims



    # Decode class integer labels

    det_labels = [rev_label_map[l] for l in det_labels[0].to('cpu').tolist()]



    # If no objects found, the detected labels will be set to ['0.'], i.e. ['background'] in SSD300.detect_objects() in model.py

    if det_labels == ['background']:

        # Just return original image

        # no objects found

        print("No object found")

        det_boxes = []

        return det_boxes



    # Annotate

    if visualize == True:

        # Color map for bounding boxes of detected objects from https://sashat.me/2017/01/11/list-of-20-simple-distinct-colors/

        distinct_colors = ['#e6194b', '#3cb44b', '#ffe119', '#0082c8', '#f58231', '#911eb4', '#46f0f0', '#f032e6',

                           '#d2f53c', '#fabebe', '#008080', '#000080', '#aa6e28', '#fffac8', '#800000', '#aaffc3',

                           '#808000',

                           '#ffd8b1', '#e6beff', '#808080', '#FFFFFF']

        label_color_map = {k: distinct_colors[i] for i, k in enumerate(labelmap.keys())}



        annotated_image = original_image

        draw = ImageDraw.Draw(annotated_image)

        font = ImageFont.load_default()



        # Suppress specific classes, if needed

        for i in range(det_boxes.size(0)):

            if suppress is not None:

                if det_labels[i] in suppress:

                    continue



            # Boxes

            box_location = det_boxes[i].tolist()

            draw.rectangle(xy=box_location, outline=label_color_map[det_labels[i]])

            # draw.rectangle(xy=[l + 1. for l in box_location], outline=label_color_map[

            #     det_labels[i]])  # a second rectangle at an offset of 1 pixel to increase line thickness

            # draw.rectangle(xy=[l + 2. for l in box_location], outline=label_color_map[

            #     det_labels[i]])  # a third rectangle at an offset of 1 pixel to increase line thickness

            # draw.rectangle(xy=[l + 3. for l in box_location], outline=label_color_map[

            #     det_labels[i]])  # a fourth rectangle at an offset of 1 pixel to increase line thickness



            # Text

            text_size = font.getsize(det_labels[i].upper())

            text_location = [box_location[0] + 2., box_location[1] - text_size[1]]

            textbox_location = [box_location[0], box_location[1] - text_size[1], box_location[0] + text_size[0] + 4.,

                                box_location[1]]

            draw.rectangle(xy=textbox_location, fill=label_color_map[det_labels[i]])

            draw.text(xy=text_location, text=det_labels[i].upper(), fill='white',

                      font=font)

        del draw

        annotated_image.save(res_path + "result_ssd.jpg")

    return det_boxes





def detect_objects(predicted_locs, predicted_scores, priors_cxcy, min_score, max_overlap, top_k):

    """

    Decipher the 8732 locations and class scores (output of ths SSD300) to detect objects.

    For each class, perform Non-Maximum Suppression (NMS) on boxes that are above a minimum threshold.

    :param predicted_locs: predicted locations/boxes w.r.t the 8732 prior boxes, a tensor of dimensions (N, 8732, 4)

    :param predicted_scores: class scores for each of the encoded locations/boxes, a tensor of dimensions (N, 8732, n_classes)

    :param min_score: minimum threshold for a box to be considered a match for a certain class

    :param max_overlap: maximum overlap two boxes can have so that the one with the lower score is not suppressed via NMS

    :param top_k: if there are a lot of resulting detection across all classes, keep only the top 'k'

    :return: detections (boxes, labels, and scores), lists of length batch_size

    """

    num_classes = opt['num_classes'] + 1  # include background



    batch_size = predicted_locs.size(0)

    n_priors = priors_cxcy.size(0)

    predicted_scores = F.softmax(predicted_scores, dim=2)  # (N, 8732, n_classes)



    # Lists to store final predicted boxes, labels, and scores for all images

    all_images_boxes = list()

    all_images_labels = list()

    all_images_scores = list()



    assert n_priors == predicted_locs.size(1) == predicted_scores.size(1)



    for i in range(batch_size):

        # Decode object coordinates from the form we regressed predicted boxes to

        decoded_locs = cxcy_to_xy(

            gcxgcy_to_cxcy(predicted_locs[i], priors_cxcy))  # (8732, 4), these are fractional pt. coordinates



        # Lists to store boxes and scores for this image

        image_boxes = list()

        image_labels = list()

        image_scores = list()



        max_scores, best_label = predicted_scores[i].max(dim=1)  # (8732)



        # Check for each class

        for c in range(1, num_classes):

            # Keep only predicted boxes and scores where scores for this class are above the minimum score

            class_scores = predicted_scores[i][:, c]  # (8732)

            score_above_min_score = class_scores > min_score  # torch.uint8 (byte) tensor, for indexing

            n_above_min_score = score_above_min_score.sum().item()

            if n_above_min_score == 0:

                continue

            class_scores = class_scores[score_above_min_score]  # (n_qualified), n_min_score <= 8732

            class_decoded_locs = decoded_locs[score_above_min_score]  # (n_qualified, 4)



            # Sort predicted boxes and scores by scores

            class_scores, sort_ind = class_scores.sort(dim=0, descending=True)  # (n_qualified), (n_min_score)

            class_decoded_locs = class_decoded_locs[sort_ind]  # (n_min_score, 4)



            # Find the overlap between predicted boxes

            overlap = find_jaccard_overlap(class_decoded_locs, class_decoded_locs)  # (n_qualified, n_min_score)



            # Non-Maximum Suppression (NMS)



            # A torch.uint8 (byte) tensor to keep track of which predicted boxes to suppress

            # 1 implies suppress, 0 implies don't suppress

            suppress = torch.zeros((n_above_min_score), dtype=torch.uint8).to(device)  # (n_qualified)



            # Consider each box in order of decreasing scores

            for box in range(class_decoded_locs.size(0)):

                # If this box is already marked for suppression

                if suppress[box] == 1:

                    continue



                # Suppress boxes whose overlaps (with this box) are greater than maximum overlap

                # Find such boxes and update suppress indices

                condition = overlap[box] > max_overlap

                condition = torch.tensor(condition, dtype=torch.uint8).to(device)

                suppress = torch.max(suppress, condition)

                # suppress = torch.max(suppress, overlap[box] > max_overlap)

                # The max operation retains previously suppressed boxes, like an 'OR' operation



                # Don't suppress this box, even though it has an overlap of 1 with itself

                suppress[box] = 0



            # Store only unsuppressed boxes for this class

            image_boxes.append(class_decoded_locs[1 - suppress])

            image_labels.append(torch.LongTensor((1 - suppress).sum().item() * [c]).to(device))

            image_scores.append(class_scores[1 - suppress])



        # If no object in any class is found, store a placeholder for 'background'

        if len(image_boxes) == 0:

            image_boxes.append(torch.FloatTensor([[0., 0., 1., 1.]]).to(device))

            image_labels.append(torch.LongTensor([0]).to(device))

            image_scores.append(torch.FloatTensor([0.]).to(device))



        # Concatenate into single tensors

        image_boxes = torch.cat(image_boxes, dim=0)  # (n_objects, 4)

        image_labels = torch.cat(image_labels, dim=0)  # (n_objects)

        image_scores = torch.cat(image_scores, dim=0)  # (n_objects)

        n_objects = image_scores.size(0)



        # Keep only the top k objects

        if n_objects > top_k:

            image_scores, sort_ind = image_scores.sort(dim=0, descending=True)

            image_scores = image_scores[:top_k]  # (top_k)

            image_boxes = image_boxes[sort_ind][:top_k]  # (top_k, 4)

            image_labels = image_labels[sort_ind][:top_k]  # (top_k)



        # Append to lists that store predicted boxes and scores for all images

        all_images_boxes.append(image_boxes)

        all_images_labels.append(image_labels)

        all_images_scores.append(image_scores)



    return all_images_boxes, all_images_labels, all_images_scores  # lists of length batch_size





if __name__ == '__main__':

    print()

    # detect_pill(image, min_score=0.2, max_overlap=0.5, top_k=3)