
# %%
import seaborn as sns

colors = sns.color_palette('colorblind')


# %%

import numpy as np

import matplotlib.pyplot as plt
from tueplots import bundles
plt.rcParams.update(bundles.icml2022())
import pandas as pd
def reward_sparsity(desired_ratio):
    return np.log(1.0 - 0.95) / np.log(desired_ratio)



ratio = np.linspace(0.0, 1.0, 100)

plt.close('all')

pr_list = []
s_list = []
data = np.zeros((100*5, 3))
for i, d in enumerate([0.05, 0.25, 0.5, 0.75, 0.95]):
    exp = reward_sparsity(d)
    s = 1.0 - ratio ** exp
    data[i*100:(i+1)*100, 0] = ratio
    data[i*100:(i+1)*100, 1] = s
    data[i*100:(i+1)*100, 2] = d


df = pd.DataFrame(columns=['PR', 'S', 'Desired Ratio'], data=data)




# %%


plt.close('all')
fig, ax  = plt.subplots()

for i, d in enumerate([0.05, 0.25, 0.5, 0.75, 0.95]):
    exp = reward_sparsity(d)
    s = 1.0 - ratio ** exp
    ax.plot(ratio, s , lw=2, label=f'{d}', color=colors[i])
    ax.vlines(x=d, ymin=0.0, ymax=0.95, color=colors[i], linestyle='--', lw=2)



ax.set_xlim((0.0, 1.0))
ax.set_ylim((0.0, 1.05))

ax.set_xlabel(r'Nodes/Edges Ratio (PR)')
ax.set_ylabel(r'Sparsity Reward ($\mathrm{R}_s$)')
ax.grid(True)
ax.legend()

plt.legend(title='Desired Ratio')



legend = ax.legend(bbox_to_anchor=(0.5, 1.55), loc='upper left', fontsize='large', ncol=1)
# legend = ax.legend(bbox_to_anchor=(1.05, 1), loc='upper left', fontsize='large',ncol=1)
handles, labels = ax.get_legend_handles_labels()
legend = ax.get_legend()
legend.remove()

plt.tight_layout()

plt.show()

fig.savefig('results_final/desired_ratio.pdf')

# create new figure and axes for legend
fig2, ax2 = plt.subplots()

# display legend on new axes
ax2.legend(handles, labels, title='maximum desired nodes/edges ratio (d)',
           loc='best',
           title_fontsize=16,
           fontsize=16, ncol=8)



# fit figure to legend size
fig2.tight_layout()

ax2.axis('off')

fig2.savefig('results_final/desired_ratio_legend.pdf')

plt.show()
