# AUTHOR RICCARDO BONETTO

import numpy as np 
import pandas as pd 
import tensorflow as tf 
import math
import os.path


from sklearn import preprocessing as pre
from sklearn.model_selection import StratifiedKFold
from sklearn.decomposition import PCA, KernelPCA

_batch_size = 64
_num_units = 1024
_learning_rate = 1e-5
_do_prob = 0.75
_max_steps = 5000
_doTest = True



def create_input_placeholders(num_features, batch_size = 64):
	X = tf.placeholder(dtype = tf.float32, shape = (None, num_features))
	labels = tf.placeholder(dtype = tf.int32, shape=(None))
	return X, labels


def perform_inference(input_placeholder, hidden_1, hidden_2, 
	num_features, training = True):
	with tf.name_scope('hidden_1'):
		weights = tf.Variable(tf.truncated_normal(
			[num_features, hidden_1], 
			stddev=1.0 / math.sqrt(float(num_features))),
			name='weights')
		biases = tf.Variable(tf.zeros(hidden_1), name='biases')
		layer_1 = tf.nn.relu(tf.matmul(input_placeholder, weights) + biases)
	with tf.name_scope('hidden_2'):
		weights = tf.Variable(tf.truncated_normal(
			[hidden_1, hidden_2], 
			stddev=1.0 / math.sqrt(float(hidden_1))),
			name='weights')
		biases = tf.Variable(tf.zeros(hidden_2), name='biases')
		layer_2 = tf.nn.relu(tf.matmul(layer_1, weights) + biases)
	with tf.name_scope('dropout'):
		do_placeholder = tf.placeholder_with_default(1.0, shape = ())
		dropout = tf.nn.dropout(layer_2, keep_prob = do_placeholder)
	with tf.name_scope('softmax'):
		weights = tf.Variable(
			tf.truncated_normal([hidden_2, 2],
			stddev=1.0 / math.sqrt(float(hidden_2))),
			name='weights')
		biases = tf.Variable(tf.zeros([2]),
			name='biases')
		logits = tf.matmul(dropout, weights) + biases
		smax = tf.nn.softmax(logits, name = 's_max')
	return logits, smax, do_placeholder


def loss_f(logits, labels):
	labels = tf.cast(labels, dtype = tf.int64)
	cross_entropy = tf.nn.sparse_softmax_cross_entropy_with_logits(
		labels = labels, logits = logits, name = 'loss_function')
	return tf.reduce_mean(cross_entropy, name='loss_mean')

def create_train_op(loss, learning_rate):
	tf.summary.scalar('loss', loss)
	optimizer = tf.train.AdamOptimizer(learning_rate)
	global_step = tf.Variable(0, name='global_step',
		trainable = False)

	train_op = optimizer.minimize(loss, 
		global_step = global_step)
	return train_op



def run_train(data, n_components, batch_size = 64):
	with tf.Graph().as_default():
		X, y_ = create_input_placeholders(n_components)
		logits_tr, smax_tr, do_placeholder = perform_inference(X, _num_units, _num_units, n_components)
		loss = loss_f(logits_tr, y_)
		tf.summary.scalar('loss', loss)
		train_op = create_train_op(loss, _learning_rate)
		summary = tf.summary.merge_all()
		init = tf.global_variables_initializer()
		saver = tf.train.Saver()
		session = tf.Session()
		summary_writer = tf.summary.FileWriter('./logs', session.graph)
		session.run(init)
		batch_counter = 0
		last_retrain = 1
		for step in range(_max_steps):
			start_offset = batch_size*batch_counter
			stop_offset = start_offset+batch_size
			batch_counter+=1
			if stop_offset<len(data['data']):
				feed_dict = {X: data['data'][
				start_offset:stop_offset],
				y_: data['labels'][
				start_offset:stop_offset],
				do_placeholder: 1-_do_prob}
				_, loss_value = session.run([train_op, loss],
					feed_dict = feed_dict)
				if step%500 == 0:
					print('Step %d: loss = %.2f' % (step, loss_value))
					summary_str = session.run(summary, feed_dict=feed_dict)
					summary_writer.add_summary(summary_str, step)
					summary_writer.flush()
				if step%1000 == 0:
					print('Training step:', step)
			else:
				random_train_indices = np.random.permutation(
					len(data['data']))
				data['data'] = data['data'][random_train_indices]
				data['labels'] = data['labels'][random_train_indices]
				batch_counter = 0
		checkpoint_file = './logs/trained_model.ckpt'
		saver.save(session, checkpoint_file)
		return session, smax_tr, train_op, loss, X, y_


def predict(data, session, input_placeholder, label_placeholder, smax, num_features):
	feed_dict = {input_placeholder: data, label_placeholder: []}
	prediction = session.run(smax, feed_dict = feed_dict)
	return np.argmax(prediction)	 
def predict_proba(data, session, input_placeholder, label_placeholder, smax, num_features):
	feed_dict = {input_placeholder: data, label_placeholder: []}
	prediction = session.run(smax, feed_dict = feed_dict)
	return prediction	 










	