# -*- coding: utf-8 -*-
"""
 @File    : grad_cam_retinanet.py
 @Time    : 2020/5/16 下午9:41
 @Author  : yizuotian
 @Description    : retinanet GradCAM
"""
import cv2
import numpy as np


class GradCAM(object):
    """
    1: 网络不更新梯度,输入需要梯度更新
    2: 使用目标类别的得分做反向传播
    """

    def __init__(self, net, layer_name):
        self.net = net
        self.layer_name = layer_name
        self.feature = []
        self.gradient = []
        self.net.eval()
        self.handlers = []
        self._register_hook()

    def _get_features_hook(self, module, input, output):
        self.feature.append(output)  # 不同层级特征
        print("feature shape:{}".format(output.size()))

    def _get_grads_hook(self, module, input_grad, output_grad):
        """

        :param input_grad: tuple, input_grad[0]: None
                                   input_grad[1]: weight
                                   input_grad[2]: bias
        :param output_grad:tuple,长度为1
        :return:
        """
        self.gradient.insert(0, output_grad[0])  # 梯度的顺序反的
        print("gradient shape:{}".format(output_grad[0].size()))

    def _register_hook(self):
        for (name, module) in self.net.named_modules():
            if name == self.layer_name:
                self.handlers.append(module.register_forward_hook(self._get_features_hook))
                self.handlers.append(module.register_backward_hook(self._get_grads_hook))

    def remove_handlers(self):
        for handle in self.handlers:
            handle.remove()

    def __call__(self, inputs, index=0):
        """

        :param inputs: {"image": [C,H,W], "height": height, "width": width}
        :param index: 第几个边框
        :return:
        """
        self.net.zero_grad()
        output = self.net.predict([inputs])
        print(output)
        score = output[0]['instances'].scores[index]
        feature_level = output[0]['instances'].feature_levels[index]  # box来自第几层feature map
        score.backward()

        gradient = self.gradient[feature_level][0].cpu().data.numpy()  # [C,H,W]
        weight = np.mean(gradient, axis=(1, 2))  # [C]

        # feature_level 指feature map的层级，0去除batch维
        feature = self.feature[feature_level][0].cpu().data.numpy()  # [C,H,W]

        cam = feature * weight[:, np.newaxis, np.newaxis]  # [C,H,W]
        cam = np.sum(cam, axis=0)  # [H,W]
        cam = np.maximum(cam, 0)  # ReLU

        # 数值归一化
        cam -= np.min(cam)
        cam /= np.max(cam)
        # 缩放到输入图像尺寸
        h, w = inputs['height'], inputs['width']
        cam = cv2.resize(cam, (w, h))

        box = output[0]['instances'].pred_boxes.tensor[index].detach().numpy().astype(np.int32)
        class_id = output[0]['instances'].pred_classes[index].detach().numpy()
        return cam, box, class_id


class GradCamPlusPlus(GradCAM):
    def __init__(self, net, layer_name):
        super(GradCamPlusPlus, self).__init__(net, layer_name)

    def __call__(self, inputs, index=0):
        """

        :param inputs: {"image": [C,H,W], "height": height, "width": width}
        :param index: 第几个边框
        :return:
        """
        self.net.zero_grad()
        output = self.net.predict([inputs])
        print(output)
        score = output[0]['instances'].scores[index]
        feature_level = output[0]['instances'].feature_levels[index]  # box来自第几层feature map
        score.backward()

        gradient = self.gradient[feature_level][0].cpu().data.numpy()  # [C,H,W]
        gradient = np.maximum(gradient, 0.)  # ReLU
        indicate = np.where(gradient > 0, 1., 0.)  # 示性函数
        norm_factor = np.sum(gradient, axis=(1, 2))  # [C]归一化
        for i in range(len(norm_factor)):
            norm_factor[i] = 1. / norm_factor[i] if norm_factor[i] > 0. else 0.  # 避免除零
        alpha = indicate * norm_factor[:, np.newaxis, np.newaxis]  # [C,H,W]

        weight = np.sum(gradient * alpha, axis=(1, 2))  # [C]  alpha*ReLU(gradient)

        feature = self.feature[feature_level][0].cpu().data.numpy()  # [C,H,W]

        cam = feature * weight[:, np.newaxis, np.newaxis]  # [C,H,W]
        cam = np.sum(cam, axis=0)  # [H,W]
        # cam = np.maximum(cam, 0)  # ReLU

        # 数值归一化
        cam -= np.min(cam)
        cam /= np.max(cam)
        # 缩放到输入图像尺寸
        h, w = inputs['height'], inputs['width']
        cam = cv2.resize(cam, (w, h))

        return cam
