from ialgebra.utils.utils_interpreter import *
from ialgebra.utils.utils_data import to_numpy, to_tensor


name2identity = {
    'grad_cam': 'GradCam',
    'grad_saliency': 'GradSaliency',
    'guided_backprop_grad': 'GuidedBackpropGrad',
    'guided_backprop_smoothgrad': 'GuidedBackpropSmoothGrad',
    'mask': 'Mask',
    'smoothgrad': 'SmoothGrad'
}


class Operator(object):
    """
    *Function*: 
    operate saliency_maps to meet user demands

    *Inputs*:
    :input:
    :model:
    :int_map: saliency_map generated by identity
    :interpreter_params:  set default if no user inputs
    :operator_params:

    *Returns*:
    :opt_map: shape = [B*C*H*W], type = numpy.ndarray
    :opt_map+img (might not use): shape = [B*C*H*W], type = numpy.ndarray 
    """
    def __init__(self, identity_name=None, dataset=None,  target_layer=None, device=None):
        # parsing identity
        self.identity = identity_name
        self.dataset = dataset
        self.target_layer = target_layer
        self._identity_class = getattr(getattr(__import__("ialgebra"), "interpreters"), name2identity[self.identity])
        # self.identity_interpreter = self._identity_class(self.model, self.dataset)


    def projection(self, bx, by, model):
        identity_interpreter = self._identity_class(model, self.dataset)
        heatmap, heatmapimg = identity_interpreter(bx, by)
        return heatmap, heatmapimg



    def selection(self, bx, by, model, region):
        identity_interpreter = self._identity_class(model, self.dataset)
        # new input
        pos0, pos1, pos2, pos3 = region[0], region[1], region[2], region[3]
        img = bx
        mat = torch.zeros(img.shape)
        roi = img[:, int(pos0):int(pos1), int(pos2):int(pos3)]
        mat[:, int(pos0):int(pos1), int(pos2):int(pos3)] = roi
        mat = mat.to(device)
        heatmap, heatmapimg = identity_interpreter(mat, by)
        return heatmap, heatmapimg 



    # same_class x1, x2, model f1, 
    def join(self, bx_list, by_list, model):
        """
        *Function*: 
        operater join: compare two inputs x and x' from same class and find most informative common features

        *Inputs*:
        :2 inputs: x, x'
        :1 model: f

        *Returns*:
        :common opt_map: 
        :opt_map+img_x:  
        :opt_map+img_x': 
        """
        bx1, by1, bx2, by2 = bx_list[0], by_list[0], bx_list[1], by_list[1]
        # bx.size=(1,3,W,H); by.size=(1)
        [bx1, bx2] = [b.unsqueeze(0) if len(b.size()) == 3 else b for b in (bx1, bx2)]
        [by1, by2] = [b.unsqueeze(0) if len(b.size()) == 0 else b for b in (by1, by2)]
        identity_interpreter = self._identity_class(model, self.dataset)

        heatmap1, heatmapimg1 = identity_interpreter(bx1, by1)
        heatmap2, heatmapimg2 = identity_interpreter(bx2, by2)

        heatmap = 0.5 * (heatmap1 + heatmap2)
        
        heatmapimg1 =  heatmap + np.float32(to_numpy(bx1))
        heatmapimg1 = (heatmapimg1 / np.max(heatmapimg1)).squeeze(0)
        heatmapimg2 =  heatmap + np.float32(to_numpy(bx2))
        heatmapimg2 = (heatmapimg2 / np.max(heatmapimg2)).squeeze(0)

        return heatmap, heatmapimg1, heatmapimg2




    def antijoin(self, bx_list, by_list, model_list, model_diff= False):
        """
        *Function*: 
        1: operater anti-join: compare two inputs x and x' from different classes and find most informative and discriminative features
        2: operater anti-join: compare one input x, and two different models f1, f2 with different classes, to find most informative and discriminative features

        *Inputs*:
        :2 inputs: x, x'
        :1 model: f

        *Returns*:
        :heatmap1
        :heatmapimg1
        :heatmap2
        :heatmapimg2
        """
        bx1, by1, bx2, by2 = bx_list[0], by_list[0], bx_list[1], by_list[1]
        # bx.size=(1,3,W,H); by.size=(1)
        [bx1, bx2] = [b.unsqueeze(0) if len(b.size()) == 3 else b for b in (bx1, bx2)]
        [by1, by2] = [b.unsqueeze(0) if len(b.size()) == 0 else b for b in (by1, by2)]
        model1, model2 = model_list[0], model_list[1]
        identity_interpreter1 = self._identity_class(model1, self.dataset)
        identity_interpreter2 = self._identity_class(model2, self.dataset)

        # case1: 1 input, 2 models
        if model_diff:
            heatmap1_1, heatmapimg1_1 = identity_interpreter1(bx1, by1)  # interpreter1_cls1
            heatmap1_2, heatmapimg1_2 = identity_interpreter2(bx1, by1)  # interpreter2_cls1
            heatmap2_1, heatmapimg2_1 = identity_interpreter1(bx1, by2)  # interpreter1_cls2
            heatmap2_2, heatmapimg2_2 = identity_interpreter2(bx1, by2)  # interpreter2_cls2

        # case2: 2 inputs, 1 model
        else:
            heatmap1_1, heatmapimg1_1 = identity_interpreter1(bx1, by1)  # interpreter1_cls1_input1
            heatmap1_2, heatmapimg1_2 = identity_interpreter2(bx1, by2)  # interpreter2_cls2_input1
            heatmap2_1, heatmapimg2_1 = identity_interpreter1(bx2, by1)  # interpreter1_cls1_input2
            heatmap2_2, heatmapimg2_2 = identity_interpreter2(bx2, by2)  # interpreter2_cls2_input2

        heatmap1 = 0.5 * (heatmap1_1 + heatmap2_1)
        heatmapimg1 =  heatmap1 + np.float32(to_numpy(bx1))
        heatmapimg1 = (heatmapimg1 / np.max(heatmapimg1)).squeeze(0)

        heatmap2 = 0.5 * (heatmap1_2 + heatmap2_2)
        heatmapimg2 =  heatmap2 + np.float32(to_numpy(bx2))
        heatmapimg2 = (heatmapimg2 / np.max(heatmapimg2)).squeeze(0)

        return heatmap1, heatmapimg1, heatmap2, heatmapimg2 





