import torch
import torchvision
import matplotlib
import matplotlib.pyplot as plt
from PIL import Image
from captum.attr import GuidedGradCam, GuidedBackprop
from captum.attr import LayerActivation, LayerConductance, LayerGradCam

from data_utils import *
from image_utils import *
from captum_utils import *
import numpy as np

from visualizers import GradCam


plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

X, y, class_names = load_imagenet_val(num=5)

# FOR THIS SECTION ONLY, we need to use gradients. We introduce a new model we will use explicitly for GradCAM for this.
gc_model = torchvision.models.squeezenet1_1(pretrained=True)
gc = GradCam()

X_tensor = torch.cat([preprocess(Image.fromarray(x)) for x in X], dim=0).requires_grad_(True)
y_tensor = torch.LongTensor(y)

# Guided Back-Propagation
gbp_result = gc.guided_backprop(X_tensor,y_tensor, gc_model)

plt.figure(figsize=(24, 24))
for i in range(gbp_result.shape[0]):
    plt.subplot(1, 5, i + 1)
    img = gbp_result[i]
    img = rescale(img)
    plt.imshow(img)
    plt.title(class_names[y[i]])
    plt.axis('off')
plt.gcf().tight_layout()
plt.savefig('visualization/guided_backprop.png')



# GradCam
# GradCAM. We have given you which module(=layer) that we need to capture gradients from, which you can see in conv_module variable below
gc_model = torchvision.models.squeezenet1_1(pretrained=True)
for param in gc_model.parameters():
    param.requires_grad = True

X_tensor = torch.cat([preprocess(Image.fromarray(x)) for x in X], dim=0).requires_grad_(True)
y_tensor = torch.LongTensor(y)
gradcam_result = gc.grad_cam(X_tensor, y_tensor, gc_model)

plt.figure(figsize=(24, 24))
for i in range(gradcam_result.shape[0]):
    gradcam_val = gradcam_result[i]
    img = X[i] + (matplotlib.cm.jet(gradcam_val)[:,:,:3]*255)
    img = img / np.max(img)
    plt.subplot(1, 5, i + 1)
    plt.imshow(img)
    plt.title(class_names[y[i]])
    plt.axis('off')
plt.gcf().tight_layout()
plt.savefig('visualization/gradcam.png')


# As a final step, we can combine GradCam and Guided Backprop to get Guided GradCam.
X_tensor = torch.cat([preprocess(Image.fromarray(x)) for x in X], dim=0).requires_grad_(True)
y_tensor = torch.LongTensor(y)
gradcam_result = gc.grad_cam(X_tensor, y_tensor, gc_model)
gbp_result = gc.guided_backprop(X_tensor, y_tensor, gc_model)

plt.figure(figsize=(24, 24))
for i in range(gradcam_result.shape[0]):
    gbp_val = gbp_result[i]
    gradcam_val = np.expand_dims(gradcam_result[i], axis=2)

    # Pointwise multiplication and normalization of the gradcam and guided backprop results (2 lines)
    img = gradcam_val * gbp_val

    img = np.expand_dims(img.transpose(2, 0, 1), axis=0)
    img = np.float32(img)
    img = torch.from_numpy(img)
    img = deprocess(img)
    plt.subplot(1, 5, i + 1)
    plt.imshow(img)
    plt.title(class_names[y[i]])
    plt.axis('off')
plt.gcf().tight_layout()
plt.savefig('visualization/guided_gradcam.png')


# **************************************************************************************** #
# Captum
model = torchvision.models.squeezenet1_1(pretrained=True)

# We don't want to train the model, so tell PyTorch not to compute gradients
# with respect to model parameters.
for param in model.parameters():
    param.requires_grad = False

# Convert X and y from numpy arrays to Torch Tensors
X_tensor = torch.cat([preprocess(Image.fromarray(x)) for x in X], dim=0)
y_tensor = torch.LongTensor(y)

conv_module = model.features[12]

##############################################################################
# TODO: Compute/Visualize GuidedBackprop and Guided GradCAM as well.         #
#       visualize_attr_maps function from captum_utils.py is useful for      #
#       visualizing captum outputs                                           #
#       Use conv_module as the convolution layer for gradcam                 #
##############################################################################
# Computing Guided GradCam

ggc = GuidedGradCam(model, conv_module)
attribution_gcc = compute_attributions(ggc, X_tensor, target = y_tensor)
# print(X_tensor.shape, y_tensor.shape, attribution_gcc.shape)
visualize_attr_maps('visualization/GuidedGradCam.png', X, y, class_names, [attribution_gcc], ['Guided_Grad_Cam'])

# Computing Guided BackProp
gbp = GuidedBackprop(model)
attribution_gbp = compute_attributions(gbp, X_tensor, target = y_tensor)
visualize_attr_maps('visualization/GuidedBackpropCam.png', X, y, class_names, [attribution_gbp], ['Guided_Backprop_Cam'])
##############################################################################
#                             END OF YOUR CODE                               #
##############################################################################

# Try out different layers and see observe how the attributions change

layer = model.features[3]

# Example visualization for using layer visualizations 
# layer_act = LayerActivation(model, layer)
# layer_act_attr = compute_attributions(layer_act, X_tensor)
# layer_act_attr_sum = layer_act_attr.mean(axis=1, keepdim=True)


##############################################################################
# TODO: Visualize Individual Layer Gradcam and Layer Conductance (similar    #
# to what we did for the other captum sections, using our helper methods),   #
# but with some preprocessing calculations.                                  #
#                                                                            #
# You can refer to the LayerActivation example above and you should be       #
# using 'layer' given above for this section                                 #
#                                                                            #
# Also note that, you would need to customize your 'attr_preprocess'         #
# parameter that you send along to 'visualize_attr_maps' as the default      #
# 'attr_preprocess' is written to only to handle multi channel attributions. #
#                                                                            #
# For layer gradcam look at the usage of the parameter relu_attributions     #
##############################################################################
# Layer gradcam aggregates across all channels
from captum.attr import LayerAttribution

N, C, H, W = X_tensor.shape

LC = LayerConductance(model, layer)
LC_attr = compute_attributions(LC, X_tensor, target = y_tensor)
LC_attr_sum = LC_attr.mean(axis = 1, keepdim = True)
LC_attr_int = LayerAttribution.interpolate(LC_attr_sum, (H,W) )
LC_attr_int = LC_attr_int.repeat(1, 3, 1, 1)
visualize_attr_maps('visualization/LayerConductance.png', X, y, class_names, [LC_attr_int], ['LayerConductance'])  

    
LGC = LayerGradCam(model, layer)
LGC_attr = compute_attributions(LGC, X_tensor, target = y_tensor)
LGC_attr_sum = LGC_attr.mean(axis = 1, keepdim = True)    
LGC_attr_int = LayerAttribution.interpolate(LGC_attr_sum, (H,W))
LGC_attr_int = LGC_attr_int.repeat(1, 3, 1, 1)
visualize_attr_maps ('visualization/LayerGradCam.png', X, y, class_names, [LGC_attr_int], ['LayerGradCam'])  
##############################################################################
#                             END OF YOUR CODE                               #
##############################################################################


[(24, 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', False, 'import torchvision\n'), (25, 'visualizers.GradCam', 'GradCam', False, 'from visualizers import GradCam\n'), (28, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (33, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (42, 'matplotlib.pyplot.savefig', 'plt.savefig', True, 'import matplotlib.pyplot as plt\n'), (48, 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', False, 'import torchvision\n'), (53, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (56, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (66, 'matplotlib.pyplot.savefig', 'plt.savefig', True, 'import matplotlib.pyplot as plt\n'), (71, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (75, 'matplotlib.pyplot.figure', 'plt.figure', True, 'import matplotlib.pyplot as plt\n'), (92, 'matplotlib.pyplot.savefig', 'plt.savefig', True, 'import matplotlib.pyplot as plt\n'), (97, 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', False, 'import torchvision\n'), (106, 'torch.LongTensor', 'torch.LongTensor', False, 'import torch\n'), (118, 'captum.attr.GuidedGradCam', 'GuidedGradCam', False, 'from captum.attr import GuidedGradCam, GuidedBackprop\n'), (124, 'captum.attr.GuidedBackprop', 'GuidedBackprop', False, 'from captum.attr import GuidedGradCam, GuidedBackprop\n'), (160, 'captum.attr.LayerConductance', 'LayerConductance', False, 'from captum.attr import LayerActivation, LayerConductance, LayerGradCam\n'), (163, 'captum.attr.LayerAttribution.interpolate', 'LayerAttribution.interpolate', False, 'from captum.attr import LayerAttribution\n'), (168, 'captum.attr.LayerGradCam', 'LayerGradCam', False, 'from captum.attr import LayerActivation, LayerConductance, LayerGradCam\n'), (171, 'captum.attr.LayerAttribution.interpolate', 'LayerAttribution.interpolate', False, 'from captum.attr import LayerAttribution\n'), (35, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (38, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (39, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (40, 'matplotlib.pyplot.axis', 'plt.axis', True, 'import matplotlib.pyplot as plt\n'), (61, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (62, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (63, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (64, 'matplotlib.pyplot.axis', 'plt.axis', True, 'import matplotlib.pyplot as plt\n'), (78, 'numpy.expand_dims', 'np.expand_dims', True, 'import numpy as np\n'), (84, 'numpy.float32', 'np.float32', True, 'import numpy as np\n'), (85, 'torch.from_numpy', 'torch.from_numpy', False, 'import torch\n'), (87, 'matplotlib.pyplot.subplot', 'plt.subplot', True, 'import matplotlib.pyplot as plt\n'), (88, 'matplotlib.pyplot.imshow', 'plt.imshow', True, 'import matplotlib.pyplot as plt\n'), (89, 'matplotlib.pyplot.title', 'plt.title', True, 'import matplotlib.pyplot as plt\n'), (90, 'matplotlib.pyplot.axis', 'plt.axis', True, 'import matplotlib.pyplot as plt\n'), (41, 'matplotlib.pyplot.gcf', 'plt.gcf', True, 'import matplotlib.pyplot as plt\n'), (60, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (65, 'matplotlib.pyplot.gcf', 'plt.gcf', True, 'import matplotlib.pyplot as plt\n'), (91, 'matplotlib.pyplot.gcf', 'plt.gcf', True, 'import matplotlib.pyplot as plt\n'), (105, 'PIL.Image.fromarray', 'Image.fromarray', False, 'from PIL import Image\n'), (59, 'matplotlib.cm.jet', 'matplotlib.cm.jet', False, 'import matplotlib\n'), (27, 'PIL.Image.fromarray', 'Image.fromarray', False, 'from PIL import Image\n'), (52, 'PIL.Image.fromarray', 'Image.fromarray', False, 'from PIL import Image\n'), (70, 'PIL.Image.fromarray', 'Image.fromarray', False, 'from PIL import Image\n')]