import numpy as np
import dill as pickle
from disropt.agents import Agent
from disropt.functions import Variable, SquaredNorm
from disropt.problems.problem import Problem
from disropt.utils.graph_constructor import MPIgraph
from disropt.utils.utilities import is_pos_def
from gradient_tracking import GradientTracking

# generate graph
graph = MPIgraph("random_binomial", "metropolis", p=0.3)
agent_id, in_nbrs, out_nbrs, in_weights, _ = graph.get_local_info()

# generate regressor and measurements
np.random.seed(agent_id)
theta_true = np.array([1, 2])
n_points = 5
D_i = np.ones((n_points, 2))
b_i = np.zeros((n_points, 1))

for j in range(n_points):
    D_i[j, 0] = 10*np.random.rand()
    b_i[j] = theta_true[0] * D_i[j, 0] + theta_true[1] + np.random.rand()

# declare a variable
dim = 2
x = Variable(dim)

# define the local objective function
objective_function = 0.5*SquaredNorm(D_i.T @ x - b_i)

# create local problem
problem = Problem(objective_function)

# create Agent and set problem
agent = Agent(in_nbrs, out_nbrs, in_weights=in_weights)
agent.set_problem(problem)

# set initial condition
x0 = 2*np.random.rand(dim, 1)

# create and run Algorithm
algorithm = GradientTracking(agent, x0, enable_log=True)
sequence = algorithm.run(iterations=1000, stepsize=0.001, verbose=True)

# get and print result
solution = algorithm.get_result().flatten()
print("Agent {} - result: {}".format(agent_id, solution))

# save number of agents
if agent_id == 0:
    np.save("agents.npy", graph.number_of_agents)

# save sequence and cost function
np.save("agent_{}_sequence.npy".format(agent_id), sequence)

with open('agent_{}_function.pkl'.format(agent_id), 'wb') as output:
    pickle.dump(agent.problem.objective_function, output)