''' This experiment is for simply measuring how norms change and compare with each other '''

import argparse
import os
import numpy as np
import numpy.linalg as LA
import matplotlib.pyplot as plt

def grad(A,w,b):
	n = A.shape[0]
	return 2*(A.T@A@w - A.T@b)/n

def experiment(n,N,d,lr,batch_size,iterations,seed):
	# seed
	np.random.seed(seed)	
	
	# print to stdout
	print("-----")
	print("1. data points (S) = {}, d = {}, workers (n) = {}, batch_size = {}".format(n,d,N,batch_size))
	print("-----")

	# is normalization required, I think they will already be
	A = np.random.randn(n,d).astype(np.float64)
	w_star = np.random.randn(d,1).astype(np.float64)
	b = A @ w_star 
	
	# arrays for plotting
	result = {
		"maxmin": [],
		"infinity_diff": [],
		"2norm_diff": [],
		"2norm": []
	}
	
	# initialize weights
	w = np.zeros((d,1))

	for it in range(iterations):
		# to shuffle and get the batches
		indices = np.arange(n)
		np.random.shuffle(indices)
		
		# get a random disjoint set of batch gradients and flatten them
		g,g0,g1 = grad(A,w,b), grad(A[indices[0:batch_size]],w,b[indices[0:batch_size]]),grad(A[indices[batch_size:]],w,b[indices[batch_size:]])
		g0,g1 = g0.flatten(),g1.flatten()
		
		# print stats
		if(it % 50 == 1):
			loss = LA.norm(A@w - b)**2/n
			print("iteration = {}, loss = {}".format(it,loss))

		# descend along full gradient
		w -= lr*g

		# norm related information
		result["maxmin"].append(np.max(g0)-np.min(g0))
		result["infinity_diff"].append(LA.norm(g0-g1,np.inf))
		result["2norm_diff"].append(LA.norm(g0-g1))
		result["2norm"].append(LA.norm(g0))

	return result


def main():
	parser = argparse.ArgumentParser(description='Norms')
	parser.add_argument('--lr', type=float, default=0.01, metavar='LR', help='learning rate (default: 0.01)')
	parser.add_argument('--iterations', type=int, default=100, metavar='ITER', help='iterations (default: 60)')
	parser.add_argument('--n', type=int, default=2**13, metavar='n', help='number of data points (default: 2^13 = 8192)')
	parser.add_argument('--d', type=int, default=100, metavar='d', help='dimension (default: 2^8 = 256)')	
	parser.add_argument('--save', action='store_true')

	args = parser.parse_args()
	
	# parameters for experiment
	n = args.n # number of data points
	d = args.d # dimension
	N = 2 # number of worker nodes
	batch_size = int(n/2) # batches you want to use for the compare
	iterations = args.iterations
	lr = args.lr
	
	# repeating the experiment for five seeds
	fields = ["maxmin","infinity_diff","2norm","2norm_diff"]
	seeds = [10]
	
	result_avg = {}
	for field in fields:
		result_avg[field] = np.zeros(iterations)

	for seed in seeds:
		result = experiment(n,N,d,lr,batch_size,iterations,seed)
		for field in fields:
			result_avg[field] += np.array(result[field])

	for field in fields:
		result_avg[field] /= len(seeds)

	# Plots
	iteration = list(range(iterations))
	# (1) norm and diff, both euclidean and infinity
	start_iter = max(0,iterations-200)
	plt.plot(iteration[start_iter:],result_avg["maxmin"][start_iter:],label='$max(g_0)-min(g_0)$')
	plt.plot(iteration[start_iter:],result_avg["infinity_diff"][start_iter:],label='$||g_0-g_1||_\infty$')
	plt.plot(iteration[start_iter:],result_avg["2norm_diff"][start_iter:],label='$||g_0-g_1||_2$')
	plt.plot(iteration[start_iter:],result_avg["2norm"][start_iter:],label='$||g_0||$')
	plt.xlabel('iteration')
	plt.ylabel('')
	plt.title('Regression norms: S = {}, n = {}\n d = {}, batch_size = {}'.format(n,N,d,batch_size),fontsize=15)
	plt.legend()
	if(args.save):
		if not os.path.isdir('out'):
			os.makedirs('out')
		plt.savefig('out/norms_S_{}_d_{}.pdf'.format(args.n,args.d))		  
	else:
		plt.show()
	plt.close()

	
if __name__ == '__main__':
	main()
