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, 11, 24, 60), 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', (), '', False, 'import torchvision\n'), ((25, 5, 25, 14), 'visualizers.GradCam', 'GradCam', ({}, {}), '()', False, 'from visualizers import GradCam\n'), ((28, 11, 28, 30), 'torch.LongTensor', 'torch.LongTensor', ({(28, 28, 28, 29): 'y'}, {}), '(y)', False, 'import torch\n'), ((33, 0, 33, 28), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((42, 0, 42, 48), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(42, 12, 42, 47): '"""visualization/guided_backprop.png"""'}, {}), "('visualization/guided_backprop.png')", True, 'import matplotlib.pyplot as plt\n'), ((48, 11, 48, 60), 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', (), '', False, 'import torchvision\n'), ((53, 11, 53, 30), 'torch.LongTensor', 'torch.LongTensor', ({(53, 28, 53, 29): 'y'}, {}), '(y)', False, 'import torch\n'), ((56, 0, 56, 28), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((66, 0, 66, 40), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(66, 12, 66, 39): '"""visualization/gradcam.png"""'}, {}), "('visualization/gradcam.png')", True, 'import matplotlib.pyplot as plt\n'), ((71, 11, 71, 30), 'torch.LongTensor', 'torch.LongTensor', ({(71, 28, 71, 29): 'y'}, {}), '(y)', False, 'import torch\n'), ((75, 0, 75, 28), 'matplotlib.pyplot.figure', 'plt.figure', (), '', True, 'import matplotlib.pyplot as plt\n'), ((92, 0, 92, 47), 'matplotlib.pyplot.savefig', 'plt.savefig', ({(92, 12, 92, 46): '"""visualization/guided_gradcam.png"""'}, {}), "('visualization/guided_gradcam.png')", True, 'import matplotlib.pyplot as plt\n'), ((97, 8, 97, 57), 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', (), '', False, 'import torchvision\n'), ((106, 11, 106, 30), 'torch.LongTensor', 'torch.LongTensor', ({(106, 28, 106, 29): 'y'}, {}), '(y)', False, 'import torch\n'), ((118, 6, 118, 39), 'captum.attr.GuidedGradCam', 'GuidedGradCam', ({(118, 20, 118, 25): 'model', (118, 27, 118, 38): 'conv_module'}, {}), '(model, conv_module)', False, 'from captum.attr import GuidedGradCam, GuidedBackprop\n'), ((124, 6, 124, 27), 'captum.attr.GuidedBackprop', 'GuidedBackprop', ({(124, 21, 124, 26): 'model'}, {}), '(model)', False, 'from captum.attr import GuidedGradCam, GuidedBackprop\n'), ((160, 5, 160, 35), 'captum.attr.LayerConductance', 'LayerConductance', ({(160, 22, 160, 27): 'model', (160, 29, 160, 34): 'layer'}, {}), '(model, layer)', False, 'from captum.attr import LayerActivation, LayerConductance, LayerGradCam\n'), ((163, 14, 163, 63), 'captum.attr.LayerAttribution.interpolate', 'LayerAttribution.interpolate', ({(163, 43, 163, 54): 'LC_attr_sum', (163, 56, 163, 61): '(H, W)'}, {}), '(LC_attr_sum, (H, W))', False, 'from captum.attr import LayerAttribution\n'), ((168, 6, 168, 32), 'captum.attr.LayerGradCam', 'LayerGradCam', ({(168, 19, 168, 24): 'model', (168, 26, 168, 31): 'layer'}, {}), '(model, layer)', False, 'from captum.attr import LayerActivation, LayerConductance, LayerGradCam\n'), ((171, 15, 171, 64), 'captum.attr.LayerAttribution.interpolate', 'LayerAttribution.interpolate', ({(171, 44, 171, 56): 'LGC_attr_sum', (171, 58, 171, 63): '(H, W)'}, {}), '(LGC_attr_sum, (H, W))', False, 'from captum.attr import LayerAttribution\n'), ((35, 4, 35, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(35, 16, 35, 17): '(1)', (35, 19, 35, 20): '(5)', (35, 22, 35, 27): '(i + 1)'}, {}), '(1, 5, i + 1)', True, 'import matplotlib.pyplot as plt\n'), ((38, 4, 38, 19), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(38, 15, 38, 18): 'img'}, {}), '(img)', True, 'import matplotlib.pyplot as plt\n'), ((39, 4, 39, 32), 'matplotlib.pyplot.title', 'plt.title', ({(39, 14, 39, 31): 'class_names[y[i]]'}, {}), '(class_names[y[i]])', True, 'import matplotlib.pyplot as plt\n'), ((40, 4, 40, 19), 'matplotlib.pyplot.axis', 'plt.axis', ({(40, 13, 40, 18): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((61, 4, 61, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(61, 16, 61, 17): '(1)', (61, 19, 61, 20): '(5)', (61, 22, 61, 27): '(i + 1)'}, {}), '(1, 5, i + 1)', True, 'import matplotlib.pyplot as plt\n'), ((62, 4, 62, 19), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(62, 15, 62, 18): 'img'}, {}), '(img)', True, 'import matplotlib.pyplot as plt\n'), ((63, 4, 63, 32), 'matplotlib.pyplot.title', 'plt.title', ({(63, 14, 63, 31): 'class_names[y[i]]'}, {}), '(class_names[y[i]])', True, 'import matplotlib.pyplot as plt\n'), ((64, 4, 64, 19), 'matplotlib.pyplot.axis', 'plt.axis', ({(64, 13, 64, 18): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((78, 18, 78, 59), 'numpy.expand_dims', 'np.expand_dims', (), '', True, 'import numpy as np\n'), ((84, 10, 84, 25), 'numpy.float32', 'np.float32', ({(84, 21, 84, 24): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((85, 10, 85, 31), 'torch.from_numpy', 'torch.from_numpy', ({(85, 27, 85, 30): 'img'}, {}), '(img)', False, 'import torch\n'), ((87, 4, 87, 28), 'matplotlib.pyplot.subplot', 'plt.subplot', ({(87, 16, 87, 17): '(1)', (87, 19, 87, 20): '(5)', (87, 22, 87, 27): '(i + 1)'}, {}), '(1, 5, i + 1)', True, 'import matplotlib.pyplot as plt\n'), ((88, 4, 88, 19), 'matplotlib.pyplot.imshow', 'plt.imshow', ({(88, 15, 88, 18): 'img'}, {}), '(img)', True, 'import matplotlib.pyplot as plt\n'), ((89, 4, 89, 32), 'matplotlib.pyplot.title', 'plt.title', ({(89, 14, 89, 31): 'class_names[y[i]]'}, {}), '(class_names[y[i]])', True, 'import matplotlib.pyplot as plt\n'), ((90, 4, 90, 19), 'matplotlib.pyplot.axis', 'plt.axis', ({(90, 13, 90, 18): '"""off"""'}, {}), "('off')", True, 'import matplotlib.pyplot as plt\n'), ((41, 0, 41, 9), 'matplotlib.pyplot.gcf', 'plt.gcf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((60, 16, 60, 27), 'numpy.max', 'np.max', ({(60, 23, 60, 26): 'img'}, {}), '(img)', True, 'import numpy as np\n'), ((65, 0, 65, 9), 'matplotlib.pyplot.gcf', 'plt.gcf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((91, 0, 91, 9), 'matplotlib.pyplot.gcf', 'plt.gcf', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((105, 33, 105, 51), 'PIL.Image.fromarray', 'Image.fromarray', ({(105, 49, 105, 50): 'x'}, {}), '(x)', False, 'from PIL import Image\n'), ((59, 18, 59, 48), 'matplotlib.cm.jet', 'matplotlib.cm.jet', ({(59, 36, 59, 47): 'gradcam_val'}, {}), '(gradcam_val)', False, 'import matplotlib\n'), ((27, 33, 27, 51), 'PIL.Image.fromarray', 'Image.fromarray', ({(27, 49, 27, 50): 'x'}, {}), '(x)', False, 'from PIL import Image\n'), ((52, 33, 52, 51), 'PIL.Image.fromarray', 'Image.fromarray', ({(52, 49, 52, 50): 'x'}, {}), '(x)', False, 'from PIL import Image\n'), ((70, 33, 70, 51), 'PIL.Image.fromarray', 'Image.fromarray', ({(70, 49, 70, 50): 'x'}, {}), '(x)', False, 'from PIL import Image\n')]