#!/usr/bin/env python
# coding: utf-8

# # 3.4 Integrated Gradients

# ## Tensorflow Walkthrough

# ### 1. Import Dependencies

# In[1]:


import os

from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.python.ops import nn_ops, gen_nn_ops
import matplotlib.pyplot as plt
import tensorflow as tf
import numpy as np

from models.models_3_4 import MNIST_CNN

get_ipython().run_line_magic('matplotlib', 'inline')

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

images = mnist.train.images
labels = mnist.train.labels

logdir = './tf_logs/3_4_IG/'
ckptdir = logdir + 'model'

if not os.path.exists(logdir):
    os.mkdir(logdir)


# ### 2. Building Graph

# In[2]:


with tf.name_scope('Classifier'):

    # Initialize neural network
    DNN = MNIST_CNN('CNN')

    # Setup training process
    X = tf.placeholder(tf.float32, [None, 784], name='X')
    Y = tf.placeholder(tf.float32, [None, 10], name='Y')

    activations, logits = DNN(X)
    
    tf.add_to_collection('IG', X)
    tf.add_to_collection('IG', logits)
    
    for activation in activations:
        tf.add_to_collection('IG', activation)

    cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y))

    optimizer = tf.train.AdamOptimizer().minimize(cost, var_list=DNN.vars)

    correct_prediction = tf.equal(tf.argmax(logits, 1), tf.argmax(Y, 1))
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

cost_summary = tf.summary.scalar('Cost', cost)
accuray_summary = tf.summary.scalar('Accuracy', accuracy)
summary = tf.summary.merge_all()


# ### 3. Training Network

# In[3]:


sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())

saver = tf.train.Saver()
file_writer = tf.summary.FileWriter(logdir, tf.get_default_graph())

# Hyper parameters
training_epochs = 15
batch_size = 100

for epoch in range(training_epochs):
    total_batch = int(mnist.train.num_examples / batch_size)
    avg_cost = 0
    avg_acc = 0
    
    for i in range(total_batch):
        batch_xs, batch_ys = mnist.train.next_batch(batch_size)
        
        # zero-center the images
        mean = np.mean(batch_xs, axis=0, keepdims=True)
        batch_xs = batch_xs - mean
        
        _, c, a, summary_str = sess.run([optimizer, cost, accuracy, summary], feed_dict={X: batch_xs, Y: batch_ys})
        avg_cost += c / total_batch
        avg_acc += a / total_batch
        
        file_writer.add_summary(summary_str, epoch * total_batch + i)

    print('Epoch:', '%04d' % (epoch + 1), 'cost =', '{:.9f}'.format(avg_cost), 'accuracy =', '{:.9f}'.format(avg_acc))
    
    saver.save(sess, ckptdir)

print('Accuracy:', sess.run(accuracy, feed_dict={X: mnist.test.images, Y: mnist.test.labels}))

sess.close()


# ### 4. Restoring Graph

# In[4]:


tf.reset_default_graph()

sess = tf.InteractiveSession()

new_saver = tf.train.import_meta_graph(ckptdir + '.meta')
new_saver.restore(sess, tf.train.latest_checkpoint(logdir))

activations = tf.get_collection('IG')
X = activations[0]
output = activations[-1]

batch, _ = mnist.train.next_batch(100)
mean = np.mean(batch, axis=0, keepdims=True)
sample_imgs = [images[np.argmax(labels, axis=1) == i][1] - mean for i in range(10)]


# ### 5. Displaying Images

# In[5]:


steps = 100
results = []

for j in range(10):
    gradient = tf.gradients(output[:,j], X)[0]
    counterfs = [sample_imgs[j][0] * i / steps for i in range(steps + 1)]
    grads = sess.run(gradient, feed_dict={X: counterfs})
    integrad = np.zeros_like(sample_imgs[j])
    
    # This is same as counterfs[-1] * np.average(grads[1:], axis=0)
    # I just used the for loop to make the Riemman Sum calculation step more explicit
    for i in range(steps):
        integrad += grads[i + 1] * (counterfs[i + 1] - counterfs[i])
    
    results.append(integrad)

plt.figure(figsize=(15,15))
for i in range(5):
    plt.subplot(5, 2, 2 * i + 1)
    plt.imshow(np.reshape(results[2 * i], [28, 28]), cmap='gray', interpolation='none')
    plt.title('Digit: {}'.format(2 * i))
    plt.xticks([])
    plt.yticks([])
    plt.colorbar()
    
    plt.subplot(5, 2, 2 * i + 2)
    plt.imshow(np.reshape(results[2 * i + 1], [28, 28]), cmap='gray', interpolation='none')
    plt.title('Digit: {}'.format(2 * i + 1))
    plt.xticks([])
    plt.yticks([])
    plt.colorbar()

plt.tight_layout()


# In[ ]: