import pdb
import numpy as np
import tensorflow as tf
import time
import math
import argparse
import sys

from net.net import CNN, CNN_model
from newton_cg import newton_cg
from utilities import read_data, predict, ConfigClass, normalize_and_reshape

def parse_args():
	parser = argparse.ArgumentParser(description='Newton method on DNN')
	parser.add_argument('--C', dest='C',
					  help='regularization term, or so-called weight decay where'+\
					  		'weight_decay = lr/(C*num_of_samples) in this implementation' ,
					  default=0.01, type=float)

	# Newton method arguments
	parser.add_argument('--GNsize', dest='GNsize',
					  help='number of samples for estimating Gauss-Newton matrix',
					  default=4096, type=int)
	parser.add_argument('--iter_max', dest='iter_max',
					  help='the maximal number of Newton iterations',
					  default=100, type=int)
	parser.add_argument('--xi', dest='xi',
					  help='the tolerance in the relative stopping condition for CG',
					  default=0.1, type=float)
	parser.add_argument('--drop', dest='drop',
					  help='the drop constants for the LM method',
					  default=2/3, type=float)
	parser.add_argument('--boost', dest='boost',
					  help='the boost constants for the LM method',
					  default=3/2, type=float)
	parser.add_argument('--eta', dest='eta',
					  help='the parameter for the line search stopping condition',
					  default=0.0001, type=float)
	parser.add_argument('--CGmax', dest='CGmax',
					  help='the maximal number of CG iterations',
					  default=250, type=int)
	parser.add_argument('--lambda', dest='_lambda',
					  help='the initial lambda for the LM method',
					  default=1, type=float)

	# SGD arguments
	parser.add_argument('--epoch_max', dest='epoch',
					  help='number of training epoch',
					  default=500, type=int)
	parser.add_argument('--lr', dest='lr',
					  help='learning rate',
					  default=0.01, type=float)
	parser.add_argument('--decay', dest='lr_decay',
					  help='learning rate decay over each mini-batch update',
					  default=0, type=float)
	parser.add_argument('--momentum', dest='momentum',
					  help='momentum of learning',
					  default=0, type=float)

	# Model training arguments
	parser.add_argument('--bsize', dest='bsize',
					  help='batch size to evaluate stochastic gradient, Gv, etc. Since the sampled data \
					  for computing Gauss-Newton matrix and etc. might not fit into memeory \
					  for one time, we will split the data into several segements and average\
					  over them.',
					  default=1024, type=int)
	parser.add_argument('--net', dest='net',
					  help='classifier type',
					  default='CNN_4layers', type=str)
	parser.add_argument('--train_set', dest='train_set',
					  help='provide the directory of .mat file for training',
					  default='data/mnist-demo.mat', type=str)
	parser.add_argument('--val_set', dest='val_set',
					  help='provide the directory of .mat file for validation',
					  default=None, type=str)
	parser.add_argument('--model', dest='model_file',
					  help='model saving address',
					  default='./saved_model/model.ckpt', type=str)
	parser.add_argument('--log', dest='log_file',
					  help='log saving directory',
					  default='./running_log/logger.log', type=str)
	parser.add_argument('--screen_log_only', dest='screen_log_only',
					  help='screen printing running log instead of storing it',
					  action='store_true')
	parser.add_argument('--optim', '-optim', 
					  help='which optimizer to use: SGD, Adam or NewtonCG',
					  default='NewtonCG', type=str)
	parser.add_argument('--loss', dest='loss', 
					  help='which loss function to use: MSELoss or CrossEntropy',
					  default='MSELoss', type=str)
	parser.add_argument('--dim', dest='dim', nargs='+', help='input dimension of data,'+\
						'shape must be:  height width num_channels',
					  default=[32, 32, 3], type=int)
	parser.add_argument('--seed', dest='seed', help='a nonnegative integer for \
						reproducibility', type=int)	  
	args = parser.parse_args()
	return args

args = parse_args()

def init_model(param):
	init_ops = []
	for p in param:
		if 'kernel' in p.name:
			weight = np.random.standard_normal(p.shape)* np.sqrt(2.0 / ((np.prod(p.get_shape().as_list()[:-1]))))
			p.assign(weight)
		elif 'bias' in p.name:
			weight = np.zeros(p.shape)
			p.assign(weight)

@tf.function
def minimize(x, y, model, optimizer, loss_with_reg):
	with tf.GradientTape() as tape:
		lwr = loss_with_reg(y, model(x))
	optimizer.minimize(lwr, model.trainable_weights, tape=tape)
	return lwr

def gradient_trainer(config, loss, model, full_batch, val_batch, test_network):	
	step = 0
	lr = config.lr
	learning_rate = lambda: lr

	reg = lambda: tf.reduce_sum([tf.reduce_sum(tf.square(p)) for p in model.trainable_weights])
	reg_const = 1/(2*config.C)
	loss_with_reg = lambda y_true, y_pred: reg_const*reg() + tf.reduce_mean(tf.reduce_sum(loss(y_true, y_pred), axis=1))

	if config.optim == 'SGD':
		optimizer = tf.keras.optimizers.SGD(
					learning_rate=learning_rate, 
					momentum=config.momentum)
	elif config.optim == 'Adam':
		optimizer = tf.keras.optimizers.AdamOptimizer(learning_rate=learning_rate,
								beta1=0.9,
								beta2=0.999,
								epsilon=1e-08)

	train_inputs, train_labels = full_batch
	num_data = train_labels.shape[0]
	num_iters = math.ceil(num_data/config.bsize)

	print(config.args)
	if not config.screen_log_only:
		log_file = open(config.log_file, 'w')
		print(config.args, file=log_file)
	

	print('-------------- initializing network by methods in He et al. (2015) --------------')
	init_model(model.trainable_weights)

	total_running_time = 0.0
	best_acc = 0.0

	for epoch in range(0, args.epoch):
		
		loss_avg = 0.0
		start = time.time()

		for i in range(num_iters):

			load_time = time.time()
			idx = np.random.choice(np.arange(0, num_data), 
					size=config.bsize, replace=False)

			batch_input = train_inputs[idx]
			batch_labels = train_labels[idx]
			batch_input = np.ascontiguousarray(batch_input)
			batch_labels = np.ascontiguousarray(batch_labels)
			config.elapsed_time += time.time() - load_time

			batch_loss = minimize(batch_input, batch_labels, model, optimizer, loss_with_reg)

			# print initial loss
			if epoch == 0 and i == 0:
				tf.print('initial f (reg + avg. loss of 1st batch):', batch_loss)
				# TODO: output to log_file if not config.screen_log_only

			loss_avg = loss_avg + batch_loss
			# print log every 10% of the iterations
			if i % math.ceil(num_iters/10) == 0:
				end = time.time()
				tf.print('Epoch ', epoch, ': ', i, '/', num_iters, ' | loss ', batch_loss, ' | lr ', lr, ' | elapsed time ', end-start, sep='')
				# TODO: output to log_file if not config.screen_log_only
			
			# adjust learning rate for SGD by inverse time decay
			if args.optim != 'Adam':
				lr = config.lr/(1 + args.lr_decay*step)
			step = step + 1

		# exclude data loading time for fair comparison
		epoch_end = time.time() - config.elapsed_time
		total_running_time += epoch_end - start
		config.elapsed_time = 0.0
		
		if val_batch is None:
			tf.print('In epoch', epoch, 'train loss:', loss_avg/(i+1), '| epoch time', epoch_end-start)
			# TODO: output to log_file if not config.screen_log_only	
		else:
			if test_network == None:
				val_loss, val_acc, _ = predict(
					loss,
					model,
					test_batch=val_batch,
					bsize=config.bsize
					)
			else:
				# A separat test network part have been done...
				val_loss, val_acc, _ = predict(
					loss,
					model,
					test_batch=val_batch,
					bsize=config.bsize
					)

			tf.print('In epoch ', epoch, ' train loss: ', loss_avg/(i+1), ' | val loss: ', val_loss, 
					' | val accuracy: ', val_acc*100, ' % | epoch time ', epoch_end-start, sep='')
			# TODO: output to log_file if not config.screen_log_only
		
			if val_acc > best_acc:
				best_acc = val_acc
				model.save_weights(config.model_file)
				print('Saved best model in {}'.format(config.model_file))

	if val_batch is None:
		model.save_weights(config.model_file)
		print('Model at the last iteration saved in {}\r\n'.format(config.model_file))
		output_str = 'total running time {:.3f}s'.format(total_running_time)
	else:
		output_str = 'Final acc: {:.3f}% | best acc {:.3f}% | total running time {:.3f}s'\
			.format(val_acc*100, best_acc*100, total_running_time)
	
	print(output_str)
	if not config.screen_log_only:
		print(output_str, file=log_file)
		log_file.close()

def newton_trainer(config, loss, model, full_batch, val_batch, test_network):

	newton_solver = newton_cg(config, outputs, loss)

	print('-------------- initializing network by methods in He et al. (2015) --------------')
	param = model.trainable_weights
	init_model(param)
	newton_solver.newton(full_batch, val_batch, saver, network, test_network)


def main():

	full_batch, num_cls, label_enum = read_data(filename=args.train_set, dim=args.dim)
	
	if args.val_set is None:
		print('No validation set is provided. Will output model at the last iteration.')
		val_batch = None
	else:
		val_batch, _, _ = read_data(filename=args.val_set, dim=args.dim, label_enum=label_enum)

	num_data = full_batch[0].shape[0]
	
	config = ConfigClass(args, num_data, num_cls)

	if isinstance(config.seed, int):
		tf.random.set_seed(config.seed)
		np.random.seed(config.seed)

	if config.net in ('CNN_4layers', 'CNN_7layers', 'VGG11', 'VGG13', 'VGG16','VGG19'):
		model = CNN_model(config.net, config.dim, config.num_cls)
	else:
		raise ValueError('Unrecognized training model')

	if config.loss == 'MSELoss':
		loss = lambda y_true, y_pred: tf.square(y_true - y_pred)
	else:
		loss = lambda y_true, y_pred: tf.nn.softmax_cross_entropy_with_logits(logits=y_pred, labels=y_true)

	full_batch[0], mean_tr = normalize_and_reshape(full_batch[0], dim=config.dim, mean_tr=None)
	if val_batch is not None:
		val_batch[0], _ = normalize_and_reshape(val_batch[0], dim=config.dim, mean_tr=mean_tr)

	param = model.trainable_weights

	# TODO: check what use_resource in tf1 means
	# TODO: check if the following variables are used, saver saves mean_param
	mean_param = tf.Variable(name='mean_tr', initial_value=mean_tr, trainable=False, 
				validate_shape=True)
	label_enum_var=tf.Variable(name='label_enum', initial_value=label_enum, trainable=False,
				validate_shape=True)
	
	if config.optim in ('SGD', 'Adam'):
		gradient_trainer(
			config, loss, model, full_batch, val_batch, test_network=None)
	elif config.optim == 'NewtonCG':
		newton_trainer(
			config, loss, model, full_batch, val_batch, test_network=None)


if __name__ == '__main__':
	main()

