import numpy as np
from disropt.algorithms import Algorithm

class GradientTracking(Algorithm):
	
	def __init__(self, agent, initial_condition, enable_log=False):
		super(GradientTracking, self).__init__(agent, enable_log)
		
		# initial condition
		self.x = initial_condition
		
		# initialize the tracker
		self.d = self.agent.problem.objective_function.subgradient(self.x)
	
	def iterate_run(self, stepsize):
		
		# communicate with neighbors
		x_neigh = self.agent.neighbors_exchange(self.x)
		d_neigh = self.agent.neighbors_exchange(self.d)
		
		# perform consensus
		W = self.agent.in_weights
		
		x_new = W[self.agent.id] * self.x
		d_new = W[self.agent.id] * self.d
		
		for j in self.agent.in_neighbors:
			x_new += W[j] * x_neigh[j]
			d_new += W[j] * d_neigh[j]
		
		# perform descent
		x_new -= stepsize * self.d
		
		# update tracker
		d_new += self.agent.problem.objective_function.subgradient(x_new) - self.agent.problem.objective_function.subgradient(self.x)
		
		# store new vectors
		self.x = x_new
		self.d = d_new
	
	def run(self, iterations=100, stepsize=0.1, verbose=False):
	
		# allocate space to store sequence
		if self.enable_log:
			dims = (iterations,) + self.x.shape # extend x.shape from the left
			sequence = np.zeros(dims)
		
		# main loop
		for k in range(iterations):
		
			# compute current stepsize
			current_stepsize = stepsize
			
			# perform one iteration
			self.iterate_run(current_stepsize)
			
			# store current solution
			if self.enable_log:
				sequence[k] = self.x
			
			# display information
			if self.agent.id == 0 and verbose:
				print('Iteration {}'.format(k), end="\r")
		
		# return solution sequence if requested
		if self.enable_log:
			return sequence
	
	def get_result(self):
		return self.x