def gioco_delle_tre_carte(df):
    df['95_mean'] = df['95']['mean']
    df['95_std'] = df['95']['std']
    return df

def pennella(df, ax, color, multiplicator):
    # Define upper and lower bounds for shaded variation
    lower_bound = df['95_mean'] + multiplicator * df['95_std'] * - 1
    upper_bound = df['95_mean'] + multiplicator * df['95_std']
    ax.fill_between(df['mean'], lower_bound, upper_bound, facecolor=color, alpha=0.2)


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

# Sparsified best combination of k and learning rate for each lambda

spar_mem_group = sparsified_mem.groupby(by=['mean', 'learning_rate', 'take'], as_index=False)
spar_mem_group = spar_mem_group.agg({'95': ['mean', 'std']}).reset_index()
# gioco delle tre carte
spar_mem_group = gioco_delle_tre_carte(spar_mem_group)
spar_mem_group = spar_mem_group.drop(['95'], axis=1)

best_spar_mem = spar_mem_group.loc[spar_mem_group.reset_index().groupby(['mean'])['95_mean'].idxmin()]
best_spar_mem
#spar_mem_group


# Sparsified best combination of k and learning rate for each lambda

spar_nomem_group = sparsified_nomem.groupby(by=['mean', 'learning_rate', 'take'], as_index=False)
spar_nomem_group = spar_nomem_group.agg({'95': ['mean', 'std']}).reset_index()
# gioco delle tre carte
spar_nomem_group = gioco_delle_tre_carte(spar_nomem_group)
spar_nomem_group = spar_nomem_group.drop(['95'], axis=1)

best_spar_nomem = spar_nomem_group.loc[spar_nomem_group.reset_index().groupby(['mean'])['95_mean'].idxmin()]
best_spar_nomem
#spar_mem_group


# Vanilla best learning rate
vanilla = cleaned_df[cleaned_df['configuration']=='AsyncPS']
vanilla = vanilla.drop([1017,2200])
# a=vanilla[vanilla['mean']==10]
# a=a[a['95']==100000]
# display(a)
group_lr_vanilla = vanilla.groupby(by=['learning_rate'], as_index=False)['95'].mean()
best_lr = group_lr_vanilla.loc[group_lr_vanilla.reset_index()['95'].idxmin()][0]
print(best_lr)
best_lr_vanilla = vanilla[vanilla['learning_rate']==best_lr]
#best_lr_vanilla = best_lr_vanilla.groupby('mean', as_index=False)['95'].mean()
best_lr_vanilla = best_lr_vanilla.groupby(by=['mean']).agg({'95': ['mean', 'std']}).reset_index()
best_lr_vanilla['95_mean'] = best_lr_vanilla['95']['mean']
best_lr_vanilla['95_std'] = best_lr_vanilla['95']['std']
best_lr_vanilla = best_lr_vanilla.drop('95', axis=1)
best_lr_vanilla



# Byzantine best learning rate for each lambda
byzantine = cleaned_df[cleaned_df['configuration']=='ByzantinePS']
group_byz = byzantine.groupby(by=['mean', 'learning_rate']).agg({'95': ['mean', 'std']}).reset_index()
#group_byz = byzantine.groupby(by=['mean', 'learning_rate']).agg({'95':'mean'}).reset_index()

group_byz['95_mean'] = group_byz['95']['mean']
group_byz['95_std'] = group_byz['95']['std']
group_byz = group_byz.drop('95', axis=1)
group_byz
#['mean']['95']['mean']
#group_byz.groupby(['mean'])['95']
best_byz = group_byz.loc[group_byz.reset_index().groupby(['mean'])['95_mean'].idxmin()]
best_byz

fig, ax = plt.subplots(figsize=(15,10))
multiplicator = (1./10.)
colors = ['green', 'blue', 'orange', 'red']

sns.lineplot(x='mean', y='95_mean', data=best_spar_mem, label=ps[0], color=colors[0])
pennella(best_spar_mem, ax, colors[0], multiplicator)

sns.lineplot(x='mean', y='95_mean', data=best_spar_nomem, label=ps[1], color=colors[1])
pennella(best_spar_nomem, ax, colors[1], multiplicator)

sns.lineplot(x='mean', y='95_mean', data=best_lr_vanilla, label=ps[2], color=colors[2])
pennella(best_lr_vanilla, ax, colors[2], multiplicator)

sns.lineplot(x='mean', y='95_mean', data=best_byz, label=ps[3], color=colors[3])
pennella(best_byz, ax, colors[3], multiplicator)
# Define upper and lower bounds for shaded variation
# lower_bound_byz = best_byz['95_mean'] + multiplicator * best_byz['95_std'] * - 1
# upper_bound_byz = best_byz['95_mean'] + multiplicator * best_byz['95_std']
# ax.fill_between(best_byz['mean'], lower_bound_byz, upper_bound_byz, facecolor='red', alpha=0.5)

# sns.lineplot(x='mean', y='90', data=best_local, label=ps[4])
# plt.yscale('log')
plt.xscale('log')
plt.xlabel('lambda')
plt.ylabel('# updates')
plt.title('# Updates to reach 95% accuracy')
plt.show()