from fully_mixed_games_utils import *
import matplotlib.pyplot as plt
import pandas as pd


figures_path = 'figures/'
########################################## PARAMETERS ##############################################
# Initializing true parameters of a fully mixed game
A1, A2, B1, B2, C1, C2, D1, D2 = 2, -1, -1, 1, -1, 3, 1, -1
A1_manip, A2_manip, B1_manip, B2_manip, C1_manip, C2_manip, D1_manip, D2_manip = A1, A2, B1, B2, C1, C2, D1, D2
# asserting conditions of a fully-mixed game
assert (A1 > C1 and D1 > B1 and B2 > A2 and C2 > D2) or (A1 < C1 and D1 < B1 and B2 < A2 and C2 < D2)

# agent parameters:
eta = 0.01
A1_manip_vect = [2]
C2_manip_vect = [3]

# simulation parameters
total_time = 50000
start_count_t = 5000
sample_size = 1

game = create_game(A1, A2, B1, B2, C1, C2, D1, D2)
manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]

####################################################################################################

actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities, total_time=total_time)

# action distribution table
history = game_final_manipulated.get_history()[start_count_t:]
p1,p2,p3,p4,p,q = 0,0,0,0,0,0
z = 0
for a in history:
    z += 1
    if a == [0,0]:
        p1 += 1; p += 1; q += 1
    elif a == [0,1]:
        p2 += 1; p += 1
    elif a == [1,0]:
        p3 += 1; q += 1
    elif a == [1,1]:
        p4 += 1
p1 /= z; p2 /= z; p3 /= z; p4 /= z; p /= z; q /= z

df = pd.DataFrame({' ':['$p = ' + str(round(2/3,3)) + '$', '$1-p = ' + str(round(1/3,3)) + '$'],
                       '$q = ' + str(round(2/5,3)) + '$': ['$'+str(round(2/3*2/5,3)) + '$', '$'+str(round(2/3*3/5,3)) + '$'],
                       '$1-q = ' + str(round(3/5,3)) + '$': ['$'+str(round(1/3*2/5,3)) + '$', '$'+str(round(1/3*3/5,3)) + '$']})
print(df.to_latex(index=False))
print('*** Empirical latex ***')
df = pd.DataFrame({' ':['$p = ' + str(round(p,3)) + '$', '$1-p = ' + str(round(1-p,3)) + '$'],
                       '$q = ' + str(round(q,3)) + '$': ['$' + str(round(p1,3))+'$', '$' + str(round(p2,3)) + '$'],
                       '$1-q = ' + str(round(1-q,3)) + '$': ['$'+str(round(p3,3))+'$', '$' + str(round(p4,3)) + '$']})
print(df.to_latex(index=False))

# plot p,q dynamics vs. time
p_history = [h[0] for h in agents_final_manipulated[0].strategy_dist_history]
q_history = [h[0] for h in agents_final_manipulated[1].strategy_dist_history]
plt.figure()
plt.plot(p_history, label='$p(t)$')
plt.plot(q_history, label='$q(t)$')
plt.legend(fontsize=12)
plt.xlabel('Time (game rounds)', fontsize=13)
plt.ylabel('Strategy profile', fontsize=13)
save_str = 'fully_mixed_p_q_dynamics+100K.PNG'
plt.savefig(figures_path + save_str, dpi=300, bbox_inches="tight")
plt.show()

# plot p,q dynamics in parametric plot
fig, ax = plt.subplots(1, 1)
plt.plot(p_history,q_history,color='black',alpha=0.8,lw=0.2)
plt.xlabel('$p(t)$', fontsize=15)
plt.ylabel('$q(t)$', fontsize=15)
save_str = 'fully_mixed_p_q_parametric_eta=001_T=50K.PNG'
plt.plot(0.5,0.5,color='black',marker='.',ms=6)
ax.set_aspect('equal', 'box')
plt.savefig(figures_path + save_str, dpi=300, bbox_inches="tight")
plt.show()

# plot Nash equilibria and empirical payoffs
game_matrix_as_list = [A1, A2, B1, B2, C1, C2, D1, D2]
game_matrix = create_2x2_game_matrix(A1, A2, B1, B2, C1, C2, D1, D2)
game = create_game(A1, A2, B1, B2, C1, C2, D1, D2)
assert game.get_game_matrix_as_array().ravel().tolist() == game_matrix_as_list

print('Game matrix:')
print(game.get_game_matrix()[0])
print(game.get_game_matrix()[1])

# plot the theoretical utilities in Nash equilibria of the true and manipulated game
A1_manip, C2_manip = 3, 1/3
manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]
plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                  utilities_manipulated=[A1, A2, B1, B2, C1, C2, D1, D2],
                  label='True NE')
plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                  utilities_manipulated=manipulated_utilities,
                  label='Users NE', marker='x', color='r')

n_simulations = 1
for _ in range(n_simulations):
    # plot the empirical utilities with truthful agents
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, [A1, A2, B1, B2, C1, C2, D1, D2])
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0] ,label='(c,d)=(2,3) Empirical MW', color='gray', marker='^')
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, [A1, A2, B1, B2, C1, C2, D1, D2], agent_type='FTPL')
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0] ,label='(c,d)=(2,3) Empirical FTPL', color='gray', marker='s')

    # plot the empirical utilities equilibrium manipulation
    A1_manip, C2_manip = 3, 1/3
    manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities)
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='Users NE - Empirical MW', color='c', marker='^')
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities, agent_type='FTPL')
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='Users NE - Empirical FTPL', color='c', marker='s')

    # plot the empirical utilities with unilateral manipulation
    A1_manip, C2_manip = 1, 3
    manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities)
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') Empirical MW', color='m', marker='^')
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities, agent_type='FTPL')
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') Empirical FTPL', color='m', marker='s')
    # plot the theoretical utilities in Nash equilibria of the true and manipulated game
    plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                      utilities_manipulated=manipulated_utilities,
                      marker='o', color='m', label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') NE')

    # plot the empirical utilities with unilateral manipulation
    A1_manip, C2_manip = 2, 1
    manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities)
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') Empirical MW', color='orange', marker='^')
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities, agent_type='FTPL')
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') Empirical FTPL', color='orange', marker='s')
    # plot the theoretical utilities in Nash equilibria of the true and manipulated game
    plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                      utilities_manipulated=manipulated_utilities,
                      marker='o', color='orange', label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') NE')

    # plot the utilities with two-player manipulation not in equilibrium but Pareto optimal
    A1_manip, C2_manip = 1, 1
    manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities, sample_size=sample_size)
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') Empirical MW', color='b', marker='^')
    actual_user_payoffs_by_c, manipulated_baseline_payoffs_by_c, \
    game_final_manipulated, agents_final_manipulated = run_simulation(eta, game, manipulated_utilities, agent_type='FTPL', sample_size=sample_size)
    mean_utilities = np.mean(manipulated_baseline_payoffs_by_c, axis=1)
    plot_empirical_payoffs(mean_utilities.tolist()[0], label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') Empirical FTPL', color='b', marker='s')
    # plot the theoretical utilities in Nash equilibria of the true and manipulated game
    plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                      utilities_manipulated=manipulated_utilities,
                      marker='o', color='b', label='(c,d)=('+str(A1_manip)+','+str(C2_manip)+') NE')

# setting legend order
handles, labels = plt.gca().get_legend_handles_labels()
order = [0,1,2,3,4,5,8,6,7,11,9,10,14,12,13]
plt.legend([handles[idx] for idx in order],[labels[idx] for idx in order], fontsize=12,
           loc='center left', bbox_to_anchor=(1, 0.5))

# plot the theoretical utilities in figure front
A1_manip, C2_manip = 3, 1/3
manipulated_utilities = [A1_manip, A2, B1, B2, C1, C2_manip, D1, D2]
plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                  utilities_manipulated=[A1, A2, B1, B2, C1, C2, D1, D2],
                  label='True NE')
plot_NE_utilities(utilities=[A1, A2, B1, B2, C1, C2, D1, D2],
                  utilities_manipulated=manipulated_utilities,
                  label='Users NE', marker='x', color='r')

plt.xlabel('$u_1$', fontsize=14)
plt.ylabel('$u_2$', fontsize=14)
save_str = 'fully_mixed_payoffs.PNG'
plt.savefig(figures_path + save_str, dpi=100, bbox_inches="tight")
plt.show()
