import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np
import sys


# ## Using IG

# ========== Functions =============
def interpolate_images(baseline,
                       image,
                       alphas):
  alphas_x = alphas[:, tf.newaxis, tf.newaxis, tf.newaxis]
  baseline_x = tf.expand_dims(baseline, axis=0)
  input_x = tf.expand_dims(image, axis=0)
  delta = input_x - baseline_x
  images = baseline_x +  alphas_x * delta
  return images


def compute_gradients(model, images, target_class_idx):
  with tf.GradientTape() as tape:
    tape.watch(images)
    logits = model(images)
    # logits is of shape (m_steps+1, nb_classes) 
    # print("logits = model(images): ", logits.shape)
    # probs output should be of shape (m_steps+1, )
    probs = logits[:, target_class_idx]
    # print("probs.shape: ", probs.shape)
  return tape.gradient(probs, images)


def integral_approximation(gradients):
  # riemann_trapezoidal
  grads = (gradients[:-1] + gradients[1:]) / tf.constant(2.0)
  integrated_gradients = tf.math.reduce_mean(grads, axis=0)
  return integrated_gradients


@tf.function
def integrated_gradients(model,
                         baseline,
                         image,
                         target_class_idx,
                         m_steps=50,
                         batch_size=32):
  # 1. Generate alphas.
  alphas = tf.linspace(start=0.0, stop=1.0, num=m_steps+1)

  # Initialize TensorArray outside loop to collect gradients.    
  gradient_batches = tf.TensorArray(tf.float32, size=m_steps+1)
    
  # Iterate alphas range and batch computation for speed, memory efficiency, and scaling to larger m_steps.
  for alpha in tf.range(0, len(alphas), batch_size):
    from_ = alpha
    to = tf.minimum(from_ + batch_size, len(alphas))
    alpha_batch = alphas[from_:to]

    # 2. Generate interpolated inputs between baseline and input.
    interpolated_path_input_batch = interpolate_images(baseline=baseline,
                                                       image=image,
                                                       alphas=alpha_batch)

    # 3. Compute gradients between model outputs and interpolated inputs.
    gradient_batch = compute_gradients(model=model, images=interpolated_path_input_batch,
                                       target_class_idx=target_class_idx)
    
    # Write batch indices and gradients to extend TensorArray.
    gradient_batches = gradient_batches.scatter(tf.range(from_, to), gradient_batch)    
  
  # Stack path gradients together row-wise into single tensor.
  total_gradients = gradient_batches.stack()

  # 4. Integral approximation through averaging gradients.
  avg_gradients = integral_approximation(gradients=total_gradients)

  # 5. Scale integrated gradients with respect to input.
  integrated_gradients = (image - baseline) * avg_gradients

  return integrated_gradients

def convergence_check(model, attributions, baseline, input, target_class_idx):
  """
  Args:
    model(keras.Model): A trained model to generate predictions and inspect.
    baseline(Tensor): A 3D image tensor with the shape 
      (image_height, image_width, 3) with the same shape as the input tensor.
    input(Tensor): A 3D image tensor with the shape 
      (image_height, image_width, 3).
    target_class_idx(Tensor): An integer that corresponds to the correct 
      ImageNet class index in the model's output predictions tensor. Default 
        value is 50 steps.   
  Returns:
    (none): Prints scores and convergence delta to sys.stdout.
  """
  # Your model's prediction on the baseline tensor. Ideally, the baseline score
  # should be close to zero.
  baseline_prediction = model(tf.expand_dims(baseline, 0))
  # print("baseline_prediction: ", baseline_prediction)
  # baseline_prediction:  tf.Tensor([[2.1683295e-04 3.1699744e-04 4.6704659e-01 5.3241956e-01]], shape=(1, 4), dtype=float32)

  baseline_score = baseline_prediction[0][target_class_idx]
  # print("baseline_score: ", baseline_score)

  # Your model's prediction and score on the input tensor.
  input_prediction = model(tf.expand_dims(input, 0))
  # print("input_prediction: ", input_prediction)
  # input_prediction:  tf.Tensor([[7.4290162e-01 2.5709778e-01 6.0866233e-07 5.7874078e-10]], shape=(1, 4), dtype=float32)

  input_score = input_prediction[0][target_class_idx]
  # print("input_score: ", input_score)

  # Sum of your IG prediction attributions.
  # print("\tattributios: ", attributions)
  ig_score = tf.math.reduce_sum(attributions)
  delta = ig_score - (input_score - baseline_score)
  # print("delta: ", delta)
  try:
    # Test your IG score is <= 5% of the input minus baseline score.
    tf.debugging.assert_near(ig_score, (input_score - baseline_score), rtol=0.05)
    tf.print('Approximation accuracy within 5%.', output_stream=sys.stdout)
  except tf.errors.InvalidArgumentError:
    tf.print('Increase or decrease m_steps to increase approximation accuracy.', output_stream=sys.stdout)
  
  tf.print('Baseline score: {:.3f}'.format(baseline_score))
  tf.print('Input score: {:.3f}'.format(input_score))
  tf.print('IG score: {:.3f}'.format(ig_score))     
  tf.print('Convergence delta: {:.3f}'.format(delta))

def plot_img_attributions(model,
                          baseline,
                          image,
                          target_class_idx,
                          m_steps=50,
                          cmap=None,
                          overlay_alpha=0.4,
                          top_prob=0.0,
                          top_label="",
                          meta={}):
  # print("\n@@@@@ plot_img_attributions called @@@@@\n")

  attributions = integrated_gradients(model=model,
                                      baseline=baseline,
                                      image=image,
                                      target_class_idx=target_class_idx,
                                      m_steps=m_steps)
  # print("\n\n\tAttributions: ", attributions)

  convergence_check(model=model,
                    attributions=attributions,
                    baseline=baseline,
                    input=image,
                    target_class_idx=target_class_idx)

  # Sum of the attributions across color channels for visualization.
  # The attribution mask shape is a grayscale image with height and width
  # equal to the original image.
  attribution_mask = tf.reduce_sum(tf.math.abs(attributions), axis=-1)

  fig, axs = plt.subplots(nrows=1, ncols=3, squeeze=False, figsize=(9,4))

  file_name = meta["file_name"]
  v = meta["v"]
  position = ""
  mode = meta["mode"]

  if mode == "Sag" and v == 1:  # sag only 1 label == 1
    position = f'P{meta["position_index"]}'
  elif mode == "Axial":
    if v == 1:
      position = "Right"
    elif v == 3:
      position = "Left"
      # flip back the v=3 crops
      attribution_mask = np.fliplr(attribution_mask)
      image = np.fliplr(image)
    elif v == 2:
      position = "Center"

  # axs[0, 0].set_title('Baseline image')
  # axs[0, 0].imshow(baseline)
  # axs[0, 0].axis('off')

  axs[0, 0].set_title('Original image')
  axs[0, 0].imshow(image)
  axs[0, 0].axis('off')

  axs[0, 1].set_title('Attribution mask')
  axs[0, 1].imshow(attribution_mask, cmap=cmap)
  axs[0, 1].axis('off')

  axs[0, 2].set_title('Overlay')
  axs[0, 2].imshow(attribution_mask, cmap=cmap)
  axs[0, 2].imshow(image, alpha=overlay_alpha)
  axs[0, 2].axis('off')

  # title and png file save name
  save_name = f'{file_name}-{mode}-{position}-{top_label}-{top_prob:0.1%}'
  fig.suptitle(save_name, fontweight='bold')
  plt.tight_layout()
  # plt.show()  # this is needed to block the process
  plt.savefig(f'{meta["save_dir"]}/{save_name}.jpeg')
  # close figure by plt.close(fig), it won't be displayed
  plt.close(fig)
  return fig


def main_ig(model, img_tensor, target_class_idx, prediction, meta):
    """
    input:
        model: center, sag, or lateral model
        img_tensor: tensor of the image for IG
        target_class_idx: index of the top pred label
        prediction: array of confidence in percentage
        meta: dict of
          file_name,
          v,
          mode,

    """
    # print("\n\n======== main_ig called ============")
    # print("target_class_idx: ", target_class_idx)
    top_prob = np.max(prediction[0])
    grading = np.array(['normal', 'mild', 'moderate', 'severe'])
    top_label = grading[target_class_idx]
    # print("img_tensor: ", img_tensor.shape, img_tensor.dtype, img_tensor[0][0])
    # ============ Constants ===================
    baseline = tf.zeros(shape=(150,150,3))

    # if needs to Visualizing gradient saturation
    visualize_grad_saturation = False
    if visualize_grad_saturation:
        m_steps = 50
        alphas = tf.linspace(start=0.0, stop=1.0, num=m_steps+1) # Generate m_steps intervals for integral_approximation() below.

        interpolated_images = interpolate_images(
            baseline=baseline,
            image=img_tensor,
            alphas=alphas)

        # ### Compute Gradients

        path_gradients = compute_gradients(
            model=model,
            images=interpolated_images,
            target_class_idx=target_class_idx)
        # print("path_gradients: ", path_gradients.shape)
        # print(np.max(path_gradients), np.min(path_gradients))
        
        # Visualize the gradient saturation
        pred = model(interpolated_images)
        pred_proba = pred[:, target_class_idx]

        plt.figure(figsize=(10, 4))
        ax1 = plt.subplot(1, 2, 1)
        ax1.plot(alphas, pred_proba)
        ax1.set_title('Target class predicted probability over alpha')
        ax1.set_ylabel('model p(target class)')
        ax1.set_xlabel('alpha')
        ax1.set_ylim([0, 1])

        ax2 = plt.subplot(1, 2, 2)
        # Average across interpolation steps
        average_grads = tf.reduce_mean(path_gradients, axis=[1, 2, 3])
        # Normalize gradients to 0 to 1 scale. E.g. (x - min(x))/(max(x)-min(x))
        average_grads_norm = (average_grads-tf.math.reduce_min(average_grads))/(tf.math.reduce_max(average_grads)-tf.reduce_min(average_grads))
        ax2.plot(alphas, average_grads_norm)
        ax2.set_title('Average pixel gradients (normalized) over alpha')
        ax2.set_ylabel('Average pixel gradients')
        ax2.set_xlabel('alpha')
        ax2.set_ylim([0, 1]);
        plt.show()

    # =========== main program ================
    # ## Visualize Attributions
    _ = plot_img_attributions(model=model,
                            image=img_tensor,
                            baseline=baseline,
                            target_class_idx=target_class_idx,
                            m_steps=240,
                            cmap=plt.cm.inferno,
                            overlay_alpha=0.4,
                            top_prob=top_prob,
                            top_label=top_label,
                            meta=meta)