# CAM
import cv2
from PIL import Image
import tensorflow as tf
from tensorflow.python.keras import backend as K

def grad_cam_keras(model, im, class_select, layer, image_size, preproc_fn, alpha=0.6, filter_threshold=0.5):
    """GradCAM method for visualizing input saliency.
    Args:
      model: keras model
      im: single image (with only RGB, [H,W,C])
      class_select: class to show
      layer: layer name
      image_size: tuple of image H,W
      preproc_fn: preprocessing function
      alpha: alpha
    Returns:
      gradient-class-activation-map
    """
    H, W = image_size[0], image_size[1]
    image = im.copy()
    if len(image) != 4:
        image = image[np.newaxis, :,:,:]
    image_original = image[0].astype("uint8")
    image = preproc_fn(image.astype("float32"))

    if class_select is None:
        output = model.predict(image)
        y_c = model.output[0, output.argmax()]
        print("Class-selected by prediction: {}".format(output.argmax()))
    else:
        y_c = model.output[0, class_select]

    conv_output = model.get_layer(layer).output
    grads = K.gradients(y_c, conv_output)[0]

    gradient_function = K.function([model.input],
                                   [conv_output, grads])
    with tf.device("/gpu:0"):
        output, grads_val = gradient_function([image])
    output, grads_val = output[0, :], grads_val[0, :, :, :]

    weights = np.mean(grads_val, axis=(0, 1))
    cam = np.dot(output, weights)

    # Process CAM
    cam = cv2.resize(cam, (H, W), cv2.INTER_LINEAR)
    cam = np.maximum(cam, 0)
    cam = cam / cam.max()

    # Filter
    cam[cam<filter_threshold] = 0

    # apply colormap
    mapping = cv2.applyColorMap(np.uint8(255 * (1-cam)), cv2.COLORMAP_JET)
    mapping = np.concatenate((mapping, ((mapping.max(axis=-1) - 128 )*255*alpha)[:,:,np.newaxis]), axis = -1)

    background = Image.fromarray(image_original)
    foreground = Image.fromarray(mapping.astype('uint8'))
    background.paste(foreground, (0, 0), foreground)

    return cam, background



  ## Usage
  cam_map, raw_img = grad_cam_keras(model=model,
                                    im=x_array[idx], ## raw image array (H,W,3)
                                    class_select=None,
                                    layer="conv5_block3_out", #"post_relu",
                                    image_size=tuple(model.input.shape.as_list()[1:3]),  # (H,W)
                                    preproc_fn=preproc_fn)