import numpy as np
import matplotlib.pyplot as plt
import dill as pickle

# initialize
with open('info.pkl', 'rb') as inp:
    info = pickle.load(inp)
NN    = info['N']
iters = info['iterations']
SS    = info['n_coupling']

# load agent data
seq = {}
local_obj_func = {}
local_coup_func = {}
local_constr = {}
for i in range(NN):
    with open('agent_{}_sequence.pkl'.format(i), 'rb') as inp:
        seq[i] = pickle.load(inp)
    with open('agent_{}_objective_func.pkl'.format(i), 'rb') as inp:
        local_obj_func[i] = pickle.load(inp)
    with open('agent_{}_coupling_func.pkl'.format(i), 'rb') as inp:
        local_coup_func[i] = pickle.load(inp)
    with open('agent_{}_local_constr.pkl'.format(i), 'rb') as inp:
        local_constr[i] = pickle.load(inp)

# compute costs and coupling values
cost_seq = np.zeros(iters)
coup_seq = np.zeros((iters, SS))

for t in range(iters):
    for i in range(NN):
        cost_seq[t] += local_obj_func[i].eval(seq[i][t, :])
        coup_seq[t] += np.squeeze(local_coup_func[i].eval(seq[i][t, :]))

# plot cost
plt.figure()
plt.title('Evolution of cost')
plt.xlabel(r"iteration $k$")
plt.ylabel(r"$\sum_{i=1}^N f_i(x_i^k)$")
plt.plot(np.arange(iters), cost_seq)

# plot maximum coupling constraint value
plt.figure()
plt.title('Evolution of maximum coupling constraint value')
plt.xlabel(r"iteration $k$")
plt.ylabel(r"$\max_{s} \: \sum_{i=1}^N g_{is}(x_i^k)$")
plt.plot(np.arange(iters), np.amax(coup_seq, axis=1))

plt.show()