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                               #
##############################################################################


[((699, 748), 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', ([], {'pretrained': '(True)'}), '(pretrained=True)\n', (731, 748), False, 'import torchvision\n'), ((754, 763), 'visualizers.GradCam', 'GradCam', ([], {}), '()\n', (761, 763), False, 'from visualizers import GradCam\n'), ((870, 889), 'torch.LongTensor', 'torch.LongTensor', (['y'], {}), '(y)\n', (886, 889), False, 'import torch\n'), ((979, 1007), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(24, 24)'}), '(figsize=(24, 24))\n', (989, 1007), True, 'import matplotlib.pyplot as plt\n'), ((1219, 1267), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""visualization/guided_backprop.png"""'], {}), "('visualization/guided_backprop.png')\n", (1230, 1267), True, 'import matplotlib.pyplot as plt\n'), ((1430, 1479), 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', ([], {'pretrained': '(True)'}), '(pretrained=True)\n', (1462, 1479), False, 'import torchvision\n'), ((1653, 1672), 'torch.LongTensor', 'torch.LongTensor', (['y'], {}), '(y)\n', (1669, 1672), False, 'import torch\n'), ((1733, 1761), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(24, 24)'}), '(figsize=(24, 24))\n', (1743, 1761), True, 'import matplotlib.pyplot as plt\n'), ((2056, 2096), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""visualization/gradcam.png"""'], {}), "('visualization/gradcam.png')\n", (2067, 2096), True, 'import matplotlib.pyplot as plt\n'), ((2289, 2308), 'torch.LongTensor', 'torch.LongTensor', (['y'], {}), '(y)\n', (2305, 2308), False, 'import torch\n'), ((2431, 2459), 'matplotlib.pyplot.figure', 'plt.figure', ([], {'figsize': '(24, 24)'}), '(figsize=(24, 24))\n', (2441, 2459), True, 'import matplotlib.pyplot as plt\n'), ((2992, 3039), 'matplotlib.pyplot.savefig', 'plt.savefig', (['"""visualization/guided_gradcam.png"""'], {}), "('visualization/guided_gradcam.png')\n", (3003, 3039), True, 'import matplotlib.pyplot as plt\n'), ((3152, 3201), 'torchvision.models.squeezenet1_1', 'torchvision.models.squeezenet1_1', ([], {'pretrained': '(True)'}), '(pretrained=True)\n', (3184, 3201), False, 'import torchvision\n'), ((3519, 3538), 'torch.LongTensor', 'torch.LongTensor', (['y'], {}), '(y)\n', (3535, 3538), False, 'import torch\n'), ((4082, 4115), 'captum.attr.GuidedGradCam', 'GuidedGradCam', (['model', 'conv_module'], {}), '(model, conv_module)\n', (4095, 4115), False, 'from captum.attr import GuidedGradCam, GuidedBackprop\n'), ((4401, 4422), 'captum.attr.GuidedBackprop', 'GuidedBackprop', (['model'], {}), '(model)\n', (4415, 4422), False, 'from captum.attr import GuidedGradCam, GuidedBackprop\n'), ((6412, 6442), 'captum.attr.LayerConductance', 'LayerConductance', (['model', 'layer'], {}), '(model, layer)\n', (6428, 6442), False, 'from captum.attr import LayerActivation, LayerConductance, LayerGradCam\n'), ((6574, 6623), 'captum.attr.LayerAttribution.interpolate', 'LayerAttribution.interpolate', (['LC_attr_sum', '(H, W)'], {}), '(LC_attr_sum, (H, W))\n', (6602, 6623), False, 'from captum.attr import LayerAttribution\n'), ((6797, 6823), 'captum.attr.LayerGradCam', 'LayerGradCam', (['model', 'layer'], {}), '(model, layer)\n', (6809, 6823), False, 'from captum.attr import LayerActivation, LayerConductance, LayerGradCam\n'), ((6964, 7014), 'captum.attr.LayerAttribution.interpolate', 'LayerAttribution.interpolate', (['LGC_attr_sum', '(H, W)'], {}), '(LGC_attr_sum, (H, W))\n', (6992, 7014), False, 'from captum.attr import LayerAttribution\n'), ((1049, 1073), 'matplotlib.pyplot.subplot', 'plt.subplot', (['(1)', '(5)', '(i + 1)'], {}), '(1, 5, i + 1)\n', (1060, 1073), True, 'import matplotlib.pyplot as plt\n'), ((1125, 1140), 'matplotlib.pyplot.imshow', 'plt.imshow', (['img'], {}), '(img)\n', (1135, 1140), True, 'import matplotlib.pyplot as plt\n'), ((1145, 1173), 'matplotlib.pyplot.title', 'plt.title', (['class_names[y[i]]'], {}), '(class_names[y[i]])\n', (1154, 1173), True, 'import matplotlib.pyplot as plt\n'), ((1178, 1193), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (1186, 1193), True, 'import matplotlib.pyplot as plt\n'), ((1933, 1957), 'matplotlib.pyplot.subplot', 'plt.subplot', (['(1)', '(5)', '(i + 1)'], {}), '(1, 5, i + 1)\n', (1944, 1957), True, 'import matplotlib.pyplot as plt\n'), ((1962, 1977), 'matplotlib.pyplot.imshow', 'plt.imshow', (['img'], {}), '(img)\n', (1972, 1977), True, 'import matplotlib.pyplot as plt\n'), ((1982, 2010), 'matplotlib.pyplot.title', 'plt.title', (['class_names[y[i]]'], {}), '(class_names[y[i]])\n', (1991, 2010), True, 'import matplotlib.pyplot as plt\n'), ((2015, 2030), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (2023, 2030), True, 'import matplotlib.pyplot as plt\n'), ((2547, 2588), 'numpy.expand_dims', 'np.expand_dims', (['gradcam_result[i]'], {'axis': '(2)'}), '(gradcam_result[i], axis=2)\n', (2561, 2588), True, 'import numpy as np\n'), ((2792, 2807), 'numpy.float32', 'np.float32', (['img'], {}), '(img)\n', (2802, 2807), True, 'import numpy as np\n'), ((2818, 2839), 'torch.from_numpy', 'torch.from_numpy', (['img'], {}), '(img)\n', (2834, 2839), False, 'import torch\n'), ((2869, 2893), 'matplotlib.pyplot.subplot', 'plt.subplot', (['(1)', '(5)', '(i + 1)'], {}), '(1, 5, i + 1)\n', (2880, 2893), True, 'import matplotlib.pyplot as plt\n'), ((2898, 2913), 'matplotlib.pyplot.imshow', 'plt.imshow', (['img'], {}), '(img)\n', (2908, 2913), True, 'import matplotlib.pyplot as plt\n'), ((2918, 2946), 'matplotlib.pyplot.title', 'plt.title', (['class_names[y[i]]'], {}), '(class_names[y[i]])\n', (2927, 2946), True, 'import matplotlib.pyplot as plt\n'), ((2951, 2966), 'matplotlib.pyplot.axis', 'plt.axis', (['"""off"""'], {}), "('off')\n", (2959, 2966), True, 'import matplotlib.pyplot as plt\n'), ((1194, 1203), 'matplotlib.pyplot.gcf', 'plt.gcf', ([], {}), '()\n', (1201, 1203), True, 'import matplotlib.pyplot as plt\n'), ((1917, 1928), 'numpy.max', 'np.max', (['img'], {}), '(img)\n', (1923, 1928), True, 'import numpy as np\n'), ((2031, 2040), 'matplotlib.pyplot.gcf', 'plt.gcf', ([], {}), '()\n', (2038, 2040), True, 'import matplotlib.pyplot as plt\n'), ((2967, 2976), 'matplotlib.pyplot.gcf', 'plt.gcf', ([], {}), '()\n', (2974, 2976), True, 'import matplotlib.pyplot as plt\n'), ((3468, 3486), 'PIL.Image.fromarray', 'Image.fromarray', (['x'], {}), '(x)\n', (3483, 3486), False, 'from PIL import Image\n'), ((1857, 1887), 'matplotlib.cm.jet', 'matplotlib.cm.jet', (['gradcam_val'], {}), '(gradcam_val)\n', (1874, 1887), False, 'import matplotlib\n'), ((798, 816), 'PIL.Image.fromarray', 'Image.fromarray', (['x'], {}), '(x)\n', (813, 816), False, 'from PIL import Image\n'), ((1581, 1599), 'PIL.Image.fromarray', 'Image.fromarray', (['x'], {}), '(x)\n', (1596, 1599), False, 'from PIL import Image\n'), ((2217, 2235), 'PIL.Image.fromarray', 'Image.fromarray', (['x'], {}), '(x)\n', (2232, 2235), False, 'from PIL import Image\n')]