#%matplotlib inline
import tensorflow as tf
import pandas as pd
import numpy as np
import math
import matplotlib.pyplot as plt
import msgpack as msg
import msgpack_numpy as m
from skimage.transform import resize
from sklearn.model_selection import train_test_split
# from sklearn.color import rgb2gray
m.patch()



def load_face_files():
    all_data = [msg.load(open('./face_images/face_images{}.bin'.format(i),'rb')) for i in range(1,6)]
    images = np.vstack([i[b'images'] for i in all_data])
    gray_images = np.dot(images,np.array([0.2125,0.7154,0.0721]))
    coords = np.vstack([i[b'co-ords'] for i in all_data])
    return gray_images,coords

images,coords = load_face_files()

def plot_image(idx,n=1):
    im = images[idx]
    part_coords = np.split(coords[idx],3)
    plt.figure(figsize=(16,16),dpi=80)
    plt.subplot(n,4,1)
    plt.imshow(im,'gray')
    for i in range(2,5):
        [x,y,w,h] = part_coords[i-2]#[:4]
        # print([x,y,w,h],all([i<0 for i in [x,y,w,h]]))
        if not all([j<0 for j in [x,y,w,h]]):
            plt.subplot(n,4,i)
            plt.imshow(im[y:y+h,x:x+w],'gray')


def get_head_images(c):
    h_idx = []
    for (idx,i) in enumerate(c):
        head_coords = np.split(i,3)[0]
        if not any([j<0 for j in head_coords]):
            h_idx.append(idx)
    return h_idx

# plot_image(958)

head_idxs = get_head_images(coords)
head_images = images[head_idxs].reshape(-1,images.shape[1]*images.shape[2])
head_coords = coords[head_idxs,:4].astype(np.float32)
tr_head_imgs,te_head_imgs,tr_head_crds,te_head_crds = train_test_split(head_images,head_coords,test_size=0.33)


def create_model(input_dim,output_dim):
    g = tf.Graph()
    with g.as_default():
        learning_rate = tf.constant(0.01,name='learning_rate')
        beta = tf.constant(0.001,name='regularization_beta')
        error_upper_bound = tf.constant(1.1,name='upper_bound')
        error_lower_bound = tf.constant(0.9,name='lower_bound')
        x = tf.placeholder(tf.float32, [None,input_dim])
        y = tf.placeholder(tf.float32, [None,output_dim])
        W1 = tf.Variable(tf.random_normal([input_dim, 512],stddev=2.0/math.sqrt(input_dim)),name='layer_1_weights')
        b1 = tf.Variable(tf.random_normal([512]),name='bias_1_weights')
        W2 = tf.Variable(tf.random_normal([512, 128],stddev=2.0/math.sqrt(512)),name='layer_2_weights')
        b2 = tf.Variable(tf.random_normal([128]),name='bias_2_weights')
        W_o = tf.Variable(tf.random_normal([128, output_dim],stddev=2.0/math.sqrt(128)),name='layer_output_weights')
        b_o = tf.Variable(tf.random_normal([output_dim]),name='bias_output_weights')
        layer_1 = tf.nn.relu(tf.add(tf.matmul(x,W1),b1))
        layer_2 = tf.nn.sigmoid(tf.add(tf.matmul(layer_1,W2),b2))
        #y_ = tf.nn.softmax(tf.add(tf.matmul(layer_2,W_o),b_o))+1e-6
        y_ = tf.add(tf.matmul(layer_2,W_o),b_o)#tf.nn.relu()

        #cross_entropy = tf.reduce_mean(-tf.reduce_sum(y * tf.log(y_)))
        # cross_entropy = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y,logits=y_))
        # distance = tf.reduce_sum(tf.square(tf.subtract(y,y_)))
        distance = tf.losses.mean_squared_error(labels = y,predictions = y_)
        regularized_loss  = distance+beta*(tf.nn.l2_loss(W1)+tf.nn.l2_loss(W2)+tf.nn.l2_loss(W_o)+tf.nn.l2_loss(b1)+tf.nn.l2_loss(b2)+tf.nn.l2_loss(b_o))
        # correct_prediction = tf.equal(tf.argmax(y,1), tf.argmax(tf.nn.softmax(y_),1))
        ratio = tf.div(y,y_)
        accuracy = tf.reduce_mean(tf.cast((ratio < error_upper_bound) & (ratio > error_lower_bound), tf.float32))
        tf.summary.scalar('distance', distance)
        tf.summary.histogram('Weights1', W1)
        tf.summary.histogram('Bias1', b1)
        tf.summary.histogram('Weights2', W2)
        tf.summary.histogram('Bias2', b2)
        tf.summary.histogram('Weights_output', W_o)
        tf.summary.histogram('Bias_output', b_o)
        merged_summary = tf.summary.merge_all()
        train_step = tf.train.GradientDescentOptimizer(learning_rate).minimize(regularized_loss)
        return (g,x,y,y_,train_step,accuracy,merged_summary)

(g,x,y,y_,train_step,accuracy,merged_summary) = create_model(tr_head_imgs.shape[1],tr_head_crds.shape[1])

def batch_data(data_x,data_y,size=128):
    batch_idxs = np.random.randint(0,data_x.shape[0],size=size)
    return (data_x[batch_idxs],data_y[batch_idxs])

def train_model(g,x,y,y_,train_step,accuracy,merged_summary):
    with g.as_default():
        with tf.Session() as s:
            train_writer = tf.summary.FileWriter('./tensor_log',s.graph)
            tf.global_variables_initializer().run()
            saver = tf.train.Saver()
            # saver.restore(s, "/tmp/model.ckpt")
            for i in range(20000):
                batch_xs,batch_ys = batch_data(tr_head_imgs,tr_head_crds,10)
                s.run([train_step],feed_dict={x:batch_xs,y:batch_ys})
                if i%100 == 0:
                    t_batch_xs,t_batch_ys = batch_data(tr_head_imgs,tr_head_crds,5)
                    [summary,acc,y_vals]= s.run([merged_summary,accuracy,y_], feed_dict={x: t_batch_xs,y: t_batch_ys})
                    train_writer.add_summary(summary,i)
                    save_path = saver.save(s, "/tmp/model.ckpt")
                    print(y_vals,t_batch_ys)
                    print("Accuracy on validation set {}".format(acc))#,'saved to ',save_path)


bb = np.array([[100,200,300,400], [100,200,300,400], [100,200,300,400], [100,200,300,400]])
bb[:,1]
train_model(g,x,y,y_,train_step,accuracy,merged_summary)