################################################################################
#Alexnet
#
#
#Original Model Code from: 
#http://www.cs.toronto.edu/~guerzhoy/tf_alexnet/
#
#
################################################################################

from numpy import *
import os
#from pylab import *
import numpy as np
#import matplotlib.pyplot as plt
#import matplotlib.cbook as cbook
import time
from scipy.misc import imread
from scipy.misc import imresize
import matplotlib.image as mpimg
from scipy.ndimage import filters
import urllib
from numpy import random
from numpy.linalg import inv
from sklearn import datasets
from sklearn.model_selection import train_test_split
import tensorflow as tf
import glob
from caffe_classes import class_names

tf.reset_default_graph()

RANDOM_SEED = 42
tf.set_random_seed(RANDOM_SEED)
C = 0.25

train_x = zeros((227,227,3)).astype(float32)
train_y = zeros((1000,18))
xdim = train_x.shape[:]
ydim = train_y.shape[0]
ydict = []
for d in os.listdir('testImg/'):
    ydict.append(d)


################################################################################
#Read parametes, images, and change to BGR

#In Python 3.5, change this to:
net_data = load(open("bvlc_alexnet.npy", "rb"), encoding="latin1").item()
#net_data = load("bvlc_alexnet.npy").item()
'''
im1 = (imread("laska.png")[:,:,:3]).astype(float32)
im1 = im1 - mean(im1)
im1[:, :, 0], im1[:, :, 2] = im1[:, :, 2], im1[:, :, 0]

im2 = (imread("poodle.png")[:,:,:3]).astype(float32)
im2[:, :, 0], im2[:, :, 2] = im2[:, :, 2], im2[:, :, 0]

im3 = (imread("dog.png")[:,:,:3]).astype(float32)
im3[:, :, 0], im3[:, :, 2] = im3[:, :, 2], im3[:, :, 0]

im4 = (imread("dog2.png")[:,:,:3]).astype(float32)
im4[:, :, 0], im4[:, :, 2] = im4[:, :, 2], im4[:, :, 0]

im5 = (imread("quail227.JPEG")[:,:,:3]).astype(float32)
im5[:, :, 0], im5[:, :, 2] = im5[:, :, 2], im5[:, :, 0]

im = [im1,im2,im3,im4,im5]
'''
im = []
k=0
for filename in glob.glob('testImg/*/*.jpeg'):
    dirname = os.path.basename(os.path.dirname(filename))
    img = (imread(filename)[:,:,:3]).astype(float32)
    img[:, :, 0], img[:, :, 2] = img[:, :, 2], img[:, :, 0]
    im.append(img)
    train_y[k][ydict.index(dirname)] = 1
    k=k+1

randomd = np.load("dict.npy")

print train_y[0]
trainx = [im[t] for t in randomd[:800]]
trainy = [train_y[t] for t in randomd[:800]]
testx = [im[t] for t in randomd[801:]]
testy = [train_y[t] for t in randomd[801:]]

################################################################################

# (self.feed('data')
#         .conv(11, 11, 96, 4, 4, padding='VALID', name='conv1')
#         .lrn(2, 2e-05, 0.75, name='norm1')
#         .max_pool(3, 3, 2, 2, padding='VALID', name='pool1')
#         .conv(5, 5, 256, 1, 1, group=2, name='conv2')
#         .lrn(2, 2e-05, 0.75, name='norm2')
#         .max_pool(3, 3, 2, 2, padding='VALID', name='pool2')
#         .conv(3, 3, 384, 1, 1, name='conv3')
#         .conv(3, 3, 384, 1, 1, group=2, name='conv4')
#         .conv(3, 3, 256, 1, 1, group=2, name='conv5')
#         .fc(4096, name='fc6')
#         .fc(4096, name='fc7')
#         .fc(1000, relu=False, name='fc8')
#         .softmax(name='prob'))


def conv(input, kernel, biases, k_h, k_w, c_o, s_h, s_w,  padding="VALID", group=1):
    '''From https://github.com/ethereon/caffe-tensorflow
    '''
    c_i = input.get_shape()[-1]
    assert c_i%group==0
    assert c_o%group==0
    convolve = lambda i, k: tf.nn.conv2d(i, k, [1, s_h, s_w, 1], padding=padding)
    
    
    if group==1:
        conv = convolve(input, kernel)
    else:
        input_groups =  tf.split(input, group, 3)   #tf.split(3, group, input)
        kernel_groups = tf.split(kernel, group, 3)  #tf.split(3, group, kernel) 
        output_groups = [convolve(i, k) for i,k in zip(input_groups, kernel_groups)]
        conv = tf.concat(output_groups, 3)          #tf.concat(3, output_groups)
    return  tf.reshape(tf.nn.bias_add(conv, biases), [-1]+conv.get_shape().as_list()[1:])

def init_weights(shape):
    """ Weight initialization """
    weights = tf.random_normal(shape, stddev=0.1)
    return tf.Variable(weights)

x = tf.placeholder(tf.float32, (None,) + xdim)

with tf.Session() as sess:

    ################################################################################
    #Layer Building

    #conv1
    #conv(11, 11, 96, 4, 4, padding='VALID', name='conv1')
    k_h = 11; k_w = 11; c_o = 96; s_h = 4; s_w = 4
    conv1W = tf.Variable(net_data["conv1"][0])
    conv1b = tf.Variable(net_data["conv1"][1])
    conv1_in = conv(x, conv1W, conv1b, k_h, k_w, c_o, s_h, s_w, padding="SAME", group=1)
    conv1 = tf.nn.relu(conv1_in)

    #lrn1
    #lrn(2, 2e-05, 0.75, name='norm1')
    radius = 2; alpha = 2e-05; beta = 0.75; bias = 1.0
    lrn1 = tf.nn.local_response_normalization(conv1,
                                                  depth_radius=radius,
                                                  alpha=alpha,
                                                  beta=beta,
                                                  bias=bias)

    #maxpool1
    #max_pool(3, 3, 2, 2, padding='VALID', name='pool1')
    k_h = 3; k_w = 3; s_h = 2; s_w = 2; padding = 'VALID'
    maxpool1 = tf.nn.max_pool(lrn1, ksize=[1, k_h, k_w, 1], strides=[1, s_h, s_w, 1], padding=padding)


    #conv2
    #conv(5, 5, 256, 1, 1, group=2, name='conv2')
    k_h = 5; k_w = 5; c_o = 256; s_h = 1; s_w = 1; group = 2
    conv2W = tf.Variable(net_data["conv2"][0])
    conv2b = tf.Variable(net_data["conv2"][1])
    conv2_in = conv(maxpool1, conv2W, conv2b, k_h, k_w, c_o, s_h, s_w, padding="SAME", group=group)
    conv2 = tf.nn.relu(conv2_in)


    #lrn2
    #lrn(2, 2e-05, 0.75, name='norm2')
    radius = 2; alpha = 2e-05; beta = 0.75; bias = 1.0
    lrn2 = tf.nn.local_response_normalization(conv2,
                                                  depth_radius=radius,
                                                  alpha=alpha,
                                                  beta=beta,
                                                  bias=bias)

    #maxpool2
    #max_pool(3, 3, 2, 2, padding='VALID', name='pool2')                                                  
    k_h = 3; k_w = 3; s_h = 2; s_w = 2; padding = 'VALID'
    maxpool2 = tf.nn.max_pool(lrn2, ksize=[1, k_h, k_w, 1], strides=[1, s_h, s_w, 1], padding=padding)

    #conv3
    #conv(3, 3, 384, 1, 1, name='conv3')
    k_h = 3; k_w = 3; c_o = 384; s_h = 1; s_w = 1; group = 1
    conv3W = tf.Variable(net_data["conv3"][0])
    conv3b = tf.Variable(net_data["conv3"][1])
    conv3_in = conv(maxpool2, conv3W, conv3b, k_h, k_w, c_o, s_h, s_w, padding="SAME", group=group)
    conv3 = tf.nn.relu(conv3_in)




    #conv4
    #conv(3, 3, 384, 1, 1, group=2, name='conv4')
    k_h = 3; k_w = 3; c_o = 384; s_h = 1; s_w = 1; group = 2
    conv4W = tf.Variable(net_data["conv4"][0])
    conv4b = tf.Variable(net_data["conv4"][1])
    conv4_in = conv(conv3, conv4W, conv4b, k_h, k_w, c_o, s_h, s_w, padding="SAME", group=group)
    conv4 = tf.nn.relu(conv4_in)


    #conv5
    #conv(3, 3, 256, 1, 1, group=2, name='conv5')
    k_h = 3; k_w = 3; c_o = 256; s_h = 1; s_w = 1; group = 2
    conv5W = tf.Variable(net_data["conv5"][0])
    conv5b = tf.Variable(net_data["conv5"][1])
    conv5_in = conv(conv4, conv5W, conv5b, k_h, k_w, c_o, s_h, s_w, padding="SAME", group=group)
    conv5 = tf.nn.relu(conv5_in)

    #maxpool5
    #max_pool(3, 3, 2, 2, padding='VALID', name='pool5')
    k_h = 3; k_w = 3; s_h = 2; s_w = 2; padding = 'VALID'
    maxpool5 = tf.nn.max_pool(conv5, ksize=[1, k_h, k_w, 1], strides=[1, s_h, s_w, 1], padding=padding)

    #fc6
    #fc(4096, name='fc6')
    fc6W = tf.Variable(net_data["fc6"][0])
    fc6b = tf.Variable(net_data["fc6"][1])
    fc6 = tf.nn.relu_layer(tf.reshape(maxpool5, [-1, int(prod(maxpool5.get_shape()[1:]))]), fc6W, fc6b)

    #fc7
    #fc(4096, name='fc7')
    fc7W = tf.Variable(net_data["fc7"][0])
    fc7b = tf.Variable(net_data["fc7"][1])
    fc7 = tf.nn.relu_layer(fc6, fc7W, fc7b)

    #fc8
    #fc(1000, relu=False, name='fc8')
    fc8W = tf.Variable(net_data["fc8"][0])
    fc8b = tf.Variable(net_data["fc8"][1])
    fc8 = tf.nn.xw_plus_b(fc7, fc8W, fc8b)


    #prob
    #softmax(name='prob'))
    prob = tf.nn.softmax(fc8)

    init = tf.initialize_all_variables()
    sess.run(init)

    t = time.time()

###########################################
    hnp = fc7.eval(feed_dict = {x:trainx})


    I_L = np.eye(hnp.shape[1])

    Q = np.matmul(np.transpose(hnp), trainy)
    omega = inv(1/C*I_L + np.matmul(np.transpose(hnp),hnp))
    print omega.shape
    
    beta = np.matmul(omega, Q)
    print beta
    yhat = np.matmul(hnp, beta)  # The \varphi function
    predict = np.argmax(yhat, axis=1)

    print yhat.shape

    print predict
    train_accuracy = np.mean(np.argmax(trainy, axis=1) ==
                                 predict)
    hnp2 = fc7.eval(feed_dict={x: testx})
    yhat2 = np.matmul(hnp2, beta)
    predict2 = np.argmax(yhat2, axis=1)
    test_accuracy  = np.mean(np.argmax(testy, axis=1) ==
                                 predict2)
    print("train accuracy = %.2f%%, test accuracy = %.2f%%"
              % (100. * train_accuracy, 100. * test_accuracy))
    #print("train accuracy = %.2f%%, "
    #          % (100. * train_accuracy))

sess.close()
tf.reset_default_graph()

#
################################################################################

#Output:
'''
for input_im_ind in range(len(output)):
    inds = argsort(output)[input_im_ind,:]
    print("Image", input_im_ind)
    for i in range(2):
        print(class_names[inds[-1-i]], output[input_im_ind][inds[-1-i]])
    print(" ")

print(time.time()-t)
'''
