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

# $\newcommand{\xv}{\mathbf{x}}
# \newcommand{\Xv}{\mathbf{X}}
# \newcommand{\yv}{\mathbf{y}}
# \newcommand{\zv}{\mathbf{z}}
# \newcommand{\av}{\mathbf{a}}
# \newcommand{\Wv}{\mathbf{W}}
# \newcommand{\wv}{\mathbf{w}}
# \newcommand{\tv}{\mathbf{t}}
# \newcommand{\Tv}{\mathbf{T}}
# \newcommand{\muv}{\boldsymbol{\mu}}
# \newcommand{\sigmav}{\boldsymbol{\sigma}}
# \newcommand{\phiv}{\boldsymbol{\phi}}
# \newcommand{\Phiv}{\boldsymbol{\Phi}}
# \newcommand{\Sigmav}{\boldsymbol{\Sigma}}
# \newcommand{\Lambdav}{\boldsymbol{\Lambda}}
# \newcommand{\half}{\frac{1}{2}}
# \newcommand{\argmax}[1]{\underset{#1}{\operatorname{argmax}}}
# \newcommand{\argmin}[1]{\underset{#1}{\operatorname{argmin}}}$

# <font size=6>**Interpreting What a Neural Network Has Learned**</font>

# [Explainable Artificial Intelligence (XAI): Concepts, taxonomies, opportunities and challenges toward responsible AI](https://www.sciencedirect.com/science/article/pii/S1566253519308103), Arrieta, et al., *Information Fusion*, Volume 58, June 2020, Pages 82-115
# 
# 
# > "Given a certain audience, explainability refers to the details and reasons a model gives to make its functioning clear or easy to understand."

# Here we will examine what the hidden units in a convolutional neural network have learned.  This is most intuitive if we focus on classification problems involving images.

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import torch
import pandas as pd
import os


# In[2]:


from A6mysolution import *


# In[3]:


# for regression problem
def rmse(a, b):
    return np.sqrt(np.mean((a - b)**2))

# for classification problem
def percent_correct(a, b):
    return 100 * np.mean(a == b)

# for classification problem
def confusion_matrix(Y_classes, T):
    class_names = np.unique(T)
    table = []
    for true_class in class_names:
        row = []
        for Y_class in class_names:
            row.append(100 * np.mean(Y_classes[T == true_class] == Y_class))
        table.append(row)
    conf_matrix = pd.DataFrame(table, index=class_names, columns=class_names)
    conf_matrix.style.background_gradient(cmap='Blues').format("{:.1f}")
    print(f'Percent Correct is {percent_correct(Y_classes, T)}')
    return conf_matrix


# In[ ]:





# In[4]:


def makeImages(nEach):
    images = np.zeros((nEach * 2, 1, 20, 20))  # nSamples, nChannels, rows, columns
    radii = 3 + np.random.randint(10 - 5, size=(nEach * 2, 1))
    centers = np.zeros((nEach * 2, 2))
    for i in range(nEach * 2):
        r = radii[i, 0]
        centers[i, :] = r + 1 + np.random.randint(18 - 2 * r, size=(1, 2))
        x = int(centers[i, 0])
        y = int(centers[i, 1])
        if i < nEach:
            # squares
            images[i, 0, x - r:x + r, y + r] = 1.0
            images[i, 0, x - r:x + r, y - r] = 1.0
            images[i, 0, x - r, y - r:y + r] = 1.0
            images[i, 0, x + r, y - r:y + r + 1] = 1.0
        else:
            # diamonds
            images[i, 0, range(x - r, x), range(y, y + r)] = 1.0
            images[i, 0, range(x - r, x), range(y, y - r, -1)] = 1.0
            images[i, 0, range(x, x + r + 1), range(y + r, y - 1, -1)] = 1.0
            images[i, 0, range(x, x + r), range(y - r, y)] = 1.0
            # images += np.random.randn(*images.shape) * 0.5
        T = np.zeros((nEach * 2, 1))
        T[nEach:] = 1
    return images, T

nEach = 1000
X, T = makeImages(nEach)
X = X.reshape(X.shape[0], -1)
print(X.shape, T.shape)

Xtest, Ttest = makeImages(nEach)
Xtest = Xtest.reshape(Xtest.shape[0], -1)

plt.plot(T);


# In[5]:


plt.imshow(-X[-1, :].reshape(20, 20), cmap='gray')
plt.xticks([])
plt.yticks([])


# In[6]:


plt.figure(figsize=(10, 3))

for i in range(10):
    plt.subplot(2, 10, i + 1)
    plt.imshow(-X[i, :].reshape(20,20), cmap='gray')
    plt.xticks([])
    plt.yticks([])

    plt.subplot(2, 10, i + 11)
    plt.imshow(-X[-i, :].reshape(20,20), cmap='gray')
    plt.xticks([])
    plt.yticks([])


# In[7]:


nnet, learning_curve = train_for_classification(X, T, hidden_layers=[10],
                                                n_epochs=500, learning_rate=0.01)
plt.plot(learning_curve);


# In[8]:


nnet


# In[9]:


Y = use(nnet, X)
Ytest = use(nnet, Xtest)
Y.shape


# In[10]:


plt.subplot(2, 1, 1)
plt.plot(Y)
plt.subplot(2, 1, 2)
plt.plot(Ytest)


# In[11]:


plt.plot(np.exp(Y))


# In[12]:


Y_classes = np.argmax(Y, axis=1).reshape(-1, 1)  # To keep 2-dimensional shape
plt.plot(Y_classes, 'o', label='Predicted')
plt.plot(T + 0.1, 'o', label='Target')
plt.legend();


# In[13]:


Y_classes_test = np.argmax(Ytest, axis=1).reshape(-1, 1)  # To keep 2-dimensional shape
plt.plot(Y_classes_test, 'o', label='Predicted')
plt.plot(T + 0.1, 'o', label='Target')
plt.legend();


# In[14]:


Y.shape


# In[15]:


confusion_matrix(Y_classes_test, Ttest)


# In[16]:


def forward_all_layers(nnet, X):
    
    X = torch.from_numpy(X).float()
    Ys = [X]
    for layer in nnet:
        Ys.append(layer(Ys[-1]))
        
    Ys = [Y.detach().numpy() for Y in Ys]    
    return Ys 


# In[17]:


Y_square = forward_all_layers(nnet, X[:10, :])
Y_diamond = forward_all_layers(nnet, X[-10:, :])


# In[18]:


nnet


# In[19]:


len(Y_square)


# In[20]:


Y_square[0].shape


# In[21]:


Y_square[1].shape


# In[22]:


plt.plot(Y_square[1]);


# In[23]:


plt.plot(Y_square[2]);


# In[24]:


both = np.vstack((Y_square[2], Y_diamond[2]))


# In[25]:


plt.plot(both);


# In[26]:


plt.figure(figsize=(15, 3))
for unit in range(10):
    plt.subplot(1, 10, unit + 1)
    plt.plot(both[:, unit])
plt.tight_layout()


# In[27]:


plt.plot(both[:, 9])


# In[28]:


nnet


# In[29]:


nnet[0].parameters()


# In[30]:


list(nnet[0].parameters())


# In[31]:


W = list(nnet[0].parameters())[0]
W = W.detach().numpy()
W.shape


# In[32]:


W = W.T
W.shape


# In[ ]:


plt.plot(W);


# In[ ]:


plt.plot(W[:, 0])


# In[ ]:


plt.imshow(W[:, 0].reshape(20, 20), cmap='RdYlGn')
plt.colorbar()


# In[ ]:


plt.figure(figsize=(15, 3))

for i in range(10):
    plt.subplot(2, 10, i + 1)
    plt.imshow(W[:, i].reshape(20,20), cmap='RdYlGn')
    plt.xticks([])
    plt.yticks([])
    plt.colorbar()
    
    plt.subplot(2, 10, i + 11)
    plt.plot(both[:, i])


# In[ ]:


X.shape


# In[ ]:


plt.imshow(X[4,:].reshape(20, 20), cmap='gray')


# Let's automate these steps in a function, so we can try different numbers of hidden units and layers.

# In[33]:


nnet


# In[34]:


Wout = list(nnet[2].parameters())[0]
Wout = Wout.detach().numpy()
Wout = Wout.T
Wout.shape


# In[35]:


def run_again(hiddens):
    
    nnet, learning_curve = train_for_classification(X, T, hidden_layers=hiddens,
                                                    n_epochs=1000, learning_rate=0.01)
    plt.figure()
    plt.plot(learning_curve)

    Y_square = forward_all_layers(nnet, X[:10, :])
    Y_diamond = forward_all_layers(nnet, X[-10:, :])
    both = np.vstack((Y_square[2], Y_diamond[2]))
    
    W = list(nnet[0].parameters())[0]
    W = W.detach().numpy()
    W = W.T
    
    Wout = list(nnet[2].parameters())[0]
    Wout = Wout.detach().numpy()
    Wout = Wout.T

    plt.figure(figsize=(15, 3))
    
    n_units = hiddens[0]

    size = int(np.sqrt(X.shape[1]))
    
    for i in range(n_units):
        plt.subplot(2, n_units, i + 1)
        plt.imshow(W[:, i].reshape(size, size), cmap='RdYlGn')
        plt.colorbar()
        plt.xticks([])
        plt.yticks([])

        plt.subplot(2, n_units, i + 1 + n_units)
        plt.plot(both[:, i])
        plt.title(f'{Wout[i,0]:.1f},{Wout[i,1]:.1f}')
        
    Y = use(nnet, X)
    Y_classes = np.argmax(Y, axis=1).reshape(-1, 1)
    print(confusion_matrix(Y_classes, T))
    
    Ytest = use(nnet, Xtest)
    Y_classes_test = np.argmax(Ytest, axis=1).reshape(-1, 1)
    print(confusion_matrix(Y_classes_test, Ttest))


# In[36]:


run_again([10])


# In[ ]:





# In[37]:


if os.path.isfile('small_mnist.npz'):
    print('Reading data from \'small_mnist.npz\'.')
    small_mnist = np.load('small_mnist.npz')
else:
    import shlex
    import subprocess
    print('Downloading small_mnist.npz from CS545 site.')
    cmd = 'curl "https://www.cs.colostate.edu/~anderson/cs545/notebooks/small_mnist.npz" -o "small_mnist.npz"'
    subprocess.call(shlex.split(cmd))
    small_mnist = np.load('small_mnist.npz')


X = small_mnist['X']
T = small_mnist['T']

X.shape, T.shape


# In[38]:


plt.imshow(-X[0, :].reshape(28, 28), cmap='gray')


# Randomly partition the data into 80% for training and 20% for testing, using the following code cells.

# In[39]:


n_samples = X.shape[0]
n_train = int(n_samples * 0.6)
rows = np.arange(n_samples)
np.random.shuffle(rows)

Xtrain = X[rows[:n_train], :]
Ttrain = T[rows[:n_train], :]
Xtest = X[rows[n_train:], :]
Ttest = T[rows[n_train:], :]


# In[59]:


def run_again_mnist(hiddens):
    
    nnet, learning_curve = train_for_classification(Xtrain, Ttrain, hidden_layers=hiddens,
                                                    n_epochs=1000, learning_rate=0.01)
    plt.figure()
    plt.plot(learning_curve)

    Y_square = forward_all_layers(nnet, X[:10, :])
    Y_diamond = forward_all_layers(nnet, X[-10:, :])
    both = np.vstack((Y_square[2], Y_diamond[2]))
    
    W = list(nnet[0].parameters())[0]
    W = W.detach().numpy()
    W = W.T
    
    Wout = list(nnet[2].parameters())[0]
    Wout = Wout.detach().numpy()
    Wout = Wout.T

    plt.figure(figsize=(15, 15))
    
    n_units = hiddens[0]

    size = int(np.sqrt(X.shape[1]))
    
    n_rows = int(np.sqrt(n_units) + 1)
    for i in range(n_units):
        plt.subplot(n_rows, n_rows, i + 1)
        plt.imshow(W[:, i].reshape(size, size), cmap='RdYlGn')
        plt.colorbar()
        plt.xticks([])
        plt.yticks([])
        
    Y = use(nnet, Xtrain)
    Y_classes = np.argmax(Y, axis=1).reshape(-1, 1)
    display(confusion_matrix(Y_classes, Ttrain))
    
    Ytest = use(nnet, Xtest)
    Y_classes_test = np.argmax(Ytest, axis=1).reshape(-1, 1)
    display(confusion_matrix(Y_classes_test, Ttest))


# In[60]:


run_again_mnist([20, 20, 20])


# In[ ]: