# Written by Roy Tseng
#
# Based on:
# --------------------------------------------------------
# Copyright (c) 2017-present, Facebook, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
##############################################################################


from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import cv2
import numpy as np
import os
import pycocotools.mask as mask_util
from enum import Enum

from mask_rcnn.utils import colormap
from miscellaneous.a2d_id_to_rgb import get_rgb
import scipy.misc as misc
import mask_rcnn.utils.keypoints as keypoint_utils

# Use a non-interactive backend
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon


class LabelMode(Enum):
    ACTOR = 1,
    ACTION = 2,
    ACTOR_ACTION = 3


plt.rcParams['pdf.fonttype'] = 42  # For editing in Adobe Illustrator


_GRAY = (218, 227, 218)
_GREEN = (18, 127, 15)
_WHITE = (255, 255, 255)


def kp_connections(keypoints):
    kp_lines = [
        [keypoints.index('left_eye'), keypoints.index('right_eye')],
        [keypoints.index('left_eye'), keypoints.index('nose')],
        [keypoints.index('right_eye'), keypoints.index('nose')],
        [keypoints.index('right_eye'), keypoints.index('right_ear')],
        [keypoints.index('left_eye'), keypoints.index('left_ear')],
        [keypoints.index('right_shoulder'), keypoints.index('right_elbow')],
        [keypoints.index('right_elbow'), keypoints.index('right_wrist')],
        [keypoints.index('left_shoulder'), keypoints.index('left_elbow')],
        [keypoints.index('left_elbow'), keypoints.index('left_wrist')],
        [keypoints.index('right_hip'), keypoints.index('right_knee')],
        [keypoints.index('right_knee'), keypoints.index('right_ankle')],
        [keypoints.index('left_hip'), keypoints.index('left_knee')],
        [keypoints.index('left_knee'), keypoints.index('left_ankle')],
        [keypoints.index('right_shoulder'), keypoints.index('left_shoulder')],
        [keypoints.index('right_hip'), keypoints.index('left_hip')],
    ]
    return kp_lines


def convert_from_cls_format(cls_boxes, cls_segms, cls_keyps):
    """Convert from the class boxes/segms/keyps format generated by the testing
    code.
    """
    box_list = [b for b in cls_boxes if len(b) > 0]
    if len(box_list) > 0:
        boxes = np.concatenate(box_list)
    else:
        boxes = None
    if cls_segms is not None:
        segms = [s for slist in cls_segms for s in slist]
    else:
        segms = None
    if cls_keyps is not None:
        keyps = [k for klist in cls_keyps for k in klist]
    else:
        keyps = None
    classes = []
    for j in range(len(cls_boxes)):
        classes += [j] * len(cls_boxes[j])
    return boxes, segms, keyps, classes


def vis_bbox_opencv(img, bbox, thick=1):
    """Visualizes a bounding box."""
    (x0, y0, w, h) = bbox
    x1, y1 = int(x0 + w), int(y0 + h)
    x0, y0 = int(x0), int(y0)
    cv2.rectangle(img, (x0, y0), (x1, y1), _GREEN, thickness=thick)
    return img


def get_class_string(class_index, score, dataset, is_actor=True):
    if is_actor:
        class_text = dataset.actor_classes[class_index] if dataset is not None else \
            'id{:d}'.format(class_index)
    else:
        class_text = dataset.action_classes[class_index] if dataset is not None else \
            'id{:d}'.format(class_index)
    return class_text + ' {:0.2f}'.format(score).lstrip('0')


def vis_one_image(
        im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
        kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
        ext='jpg'):
    """Visual debugging of detections."""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if isinstance(boxes, list):
        boxes, segms, keypoints, classes = convert_from_cls_format(
            boxes, segms, keypoints)

    if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
        return

    if segms is not None:
        masks = mask_util.decode(segms)

    color_list = colormap.colormap(rgb=True) / 255

    # dataset_keypoints, _ = keypoint_utils.get_keypoints()
    # kp_lines = kp_connections(dataset_keypoints)
    # cmap = plt.get_cmap('rainbow')
    # colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]

    fig = plt.figure(frameon=False)
    fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.axis('off')
    fig.add_axes(ax)
    ax.imshow(im)

    # Display in largest to smallest order to reduce occlusion
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    sorted_inds = np.argsort(-areas)

    mask_color_id = 0
    for i in sorted_inds:
        bbox = boxes[i, :4]
        # score = boxes[i, -1]
        actor_score = boxes[i, -3]
        action_score = boxes[i, -2]
        action_cls = int(boxes[i, -1])
        if actor_score < thresh or action_score < thresh:
            continue

        # print(dataset.classes[classes[i]], score)
        # show box (off by default, box_alpha=0.0)
        ax.add_patch(
            plt.Rectangle((bbox[0], bbox[1]),
                          bbox[2] - bbox[0],
                          bbox[3] - bbox[1],
                          fill=False, edgecolor='g',
                          linewidth=0.5, alpha=box_alpha))

        if show_class:
            ax.text(
                bbox[0], bbox[1] - 2,
                get_class_string(classes[i], actor_score, dataset, is_actor=True)
                + '--' + get_class_string(action_cls, action_score, dataset, is_actor=False),
                fontsize=3,
                family='serif',
                bbox=dict(
                    facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
                color='white')

        # show mask
        if segms is not None and len(segms) > i:
            img = np.ones(im.shape)
            color_mask = color_list[mask_color_id % len(color_list), 0:3]
            mask_color_id += 1

            w_ratio = .4
            for c in range(3):
                color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
            for c in range(3):
                img[:, :, c] = color_mask[c]
            e = masks[:, :, i]

            _, contour, hier = cv2.findContours(
                e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)

            for c in contour:
                polygon = Polygon(
                    c.reshape((-1, 2)),
                    fill=True, facecolor=color_mask,
                    edgecolor='w', linewidth=1.2,
                    alpha=0.5)
                ax.add_patch(polygon)

        output_name = os.path.basename(im_name) + '.' + ext
        fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
        plt.close('all')


def save_segm_result_with_rgb(
        im, im_name, output_dir, rgb_output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
        kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False, label_mode=LabelMode.ACTOR_ACTION,
        ext='npy'):
    """Visual debugging of detections."""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    if isinstance(boxes, list):
        boxes, segms, keypoints, classes = convert_from_cls_format(
            boxes, segms, keypoints)

    output_name = os.path.basename(im_name) + '.' + ext
    rgb_output_name = os.path.basename(im_name) + '.' + 'jpg'
    fname = os.path.join(output_dir, '{}'.format(output_name))
    jpg_name = os.path.join(rgb_output_dir, '{}'.format(rgb_output_name))
    h, w, _ = im.shape
    seg_result = np.zeros(shape=(h, w))
    rgb_seg_result = np.zeros(shape=(h, w, 3))

    if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
        # no boxes detected
        np.save(fname, seg_result)
        misc.imsave(jpg_name, rgb_seg_result.astype(np.uint8))
        return

    assert segms is not None
    masks = mask_util.decode(segms)
    # dataset_keypoints, _ = keypoint_utils.get_keypoints()
    # kp_lines = kp_connections(dataset_keypoints)
    # cmap = plt.get_cmap('rainbow')
    # colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]

    # fig = plt.figure(frameon=False)
    # fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
    # ax = plt.Axes(fig, [0., 0., 1., 1.])
    # ax.axis('off')
    # fig.add_axes(ax)
    # ax.imshow(im)

    # Display in largest to smallest order to reduce occlusion
    areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
    sorted_inds = np.argsort(-areas)

    # mask_color_id = 0
    for i in sorted_inds:
        # bbox = boxes[i, :4]
        # score = boxes[i, -1]
        actor_score = boxes[i, -3]
        action_score = boxes[i, -2]
        action_cls = int(boxes[i, -1])
        if actor_score < thresh or action_score < thresh:
            continue

        non_zero_area = np.where(seg_result > 0)[0]
        actor_action_id = int(classes[i]) * 10 + action_cls
        cur_instance_mask = masks[:, :, i]
        # mask with front instances removed
        cur_instance_mask[non_zero_area] = 0
        rgb_vis = np.repeat(cur_instance_mask[:, :, np.newaxis], repeats=3, axis=2) * np.array(get_rgb(actor_action_id))
        seg_result += cur_instance_mask * actor_action_id
        rgb_seg_result += rgb_vis

    np.save(fname, seg_result)
    misc.imsave(jpg_name, rgb_seg_result.astype(np.uint8))
