import os 
import numpy as np
import tensorflow as tf
import pandas as pd  
import matplotlib.pyplot as plt  
import sys
#You have freedom of using eager execution in tensorflow
plt.rcParams['figure.figsize'] = (10.0, 8.0) # set default size of plots
plt.rcParams['image.interpolation'] = 'nearest'
plt.rcParams['image.cmap'] = 'gray'

'''

IST 597: Foundations of Deep Learning
Problem 1b: Softmax Regression \& the Spiral Problem

    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
    
        /////////******************/////////
    //// Umar Farooq Mohammad  /////////
     ////  ubm5020@psu.edu     /////////
    /////////******************/////////
'''

# NOTE: you will need to tinker with the meta-parameters below yourself (do not think of them as defaults by any means)
# NOTE: You can work in pure eager mode or Keras or even hybrid
# NOTE : Use tf.Variable when using gradientTape, if you build your own gradientTape then use simple linear algebra using numpy or tensorflow math
	
def computeGrad(X,y,theta,reg): # returns nabla
    # WRITEME: write your code here to complete the routine
    W = theta[0]
    b = theta[1]
    size = X.shape[0]
    scores,probs = predict(X,theta)
    probs = probs.numpy()
    for i in range(len(y)):
        temp = y[i]
        probs[i][y[temp]] -= 1
        
        
    dW = np.matmul(tf.transpose(X),probs) + reg*W
    db =(1/size)*probs
    db = tf.math.reduce_sum(db,axis=0)
    return (dW,db)

def computeCost(X,y,theta,reg):
    # WRITEME: write your code here to complete the routine
    W = theta[0]
    scores,probs = predict(X,theta)
    probs = probs.numpy()
    size = X.shape[0]
    probs = -np.log(probs)

    cost = 0
    for i in range(len(y)):
        temp = y[i]
        cost += probs[i][temp]

    
    return (cost/size) + ((reg/2)*np.sum(np.power(W.numpy(),2)))

def predict(X,theta):
    # WRITEME: write your code here to complete the routine
    W = theta[0]
    b = theta[1]
    # evaluate class scores

    scores = np.matmul(X,W) + b
    # compute the class probabilities
    A = tf.exp(scores)
    B = tf.math.reduce_sum(A,axis=1)
    return (scores,tf.transpose((tf.transpose(A)/B)))




np.random.seed(21131186) #Provide your unique Random seed
# Load in the data from disk
path = os.getcwd() + '/data/spiral_train.dat'  
data = pd.read_csv(path, header=None) 

# set X (training data) and y (target variable)
cols = data.shape[1]  
X = data.iloc[:,0:cols-1]  
y = data.iloc[:,cols-1:cols] 

# convert from data frames to numpy matrices
X = np.array(X.values)  
y = np.array(y.values)
y = y.flatten()

X_tf = tf.constant(X) # Or simply do X = X_tf (if using numpy) or convert to tensorflow objects
Y_tf = tf.constant(y)
# initialize parameters randomly
D = X.shape[1]
K = np.amax(y) + 1

#Train a Linear Classifier
# initialize parameters randomly
initializer = tf.random_normal_initializer(mean=0.0, stddev=0.01, seed=21131186)
W = tf.Variable(initializer([D, K]))
W = tf.cast(W,dtype = tf.double)
b = tf.Variable(tf.random.normal([K]))


# some hyperparameters
n_e = 100
check = 10 # every so many pass/epochs, print loss/error to terminal
step_size = 0.1
reg = 0.001 # regularization strength
losses = []
# gradient descent loop
for i in range(n_e):
    # WRITEME: write your code here to perform a step of gradient descent & record anything else desired for later
    theta = (W,b)
    
    loss = computeCost(X,y,theta,reg)
    losses.append(loss)
    if i % check == 0:
            print("iteration %d: loss %f" % (i, loss))

    # perform a parameter update
    # WRITEME: write your update rule(s) here
    dW,db = computeGrad(X,y,theta,reg)
    W = W - (step_size*dW)
    b = b - (step_size*db)
    
 
# TODO: remove this line below once you have correctly implemented/gradient-checked your various sub-routines
#sys.exit(0) 
  
# evaluate training set accuracy
scores, probs = predict(X,theta)
#scores = np.dot(X, W) + b
predicted_class = (tf.argmax(scores, axis=1))
print('training accuracy: %.2f' % ((np.mean(predicted_class == y))))


# loss per epoch plot

plt.title("Training and Validation cost per epoch")
plt.xlabel("epoch number")
plt.ylabel("cost for this epoch")
plt.plot(losses)

# plot the resulting classifier
h = 0.02
x_min, x_max = X[:, 0].min() - 1, X[:, 0].max() + 1
y_min, y_max = X[:, 1].min() - 1, X[:, 1].max() + 1
xx, yy = np.meshgrid(np.arange(x_min, x_max, h),
                     np.arange(y_min, y_max, h))
Z, P = predict(np.concatenate((xx.ravel().reshape(-1,1), yy.ravel().reshape(-1,1)),axis=1), theta)
Z = np.argmax(Z, axis=1)
Z = Z.reshape(xx.shape)
fig = plt.figure()
plt.contourf(xx, yy, Z, cmap=plt.cm.Spectral, alpha=0.8)
plt.scatter(X[:, 0], X[:, 1], c=y, s=40, cmap=plt.cm.Spectral)
plt.xlim(xx.min(), xx.max())
plt.ylim(yy.min(), yy.max())
#fig.savefig('spiral_linear.png')

plt.show()
