import numpy as np
from matplotlib import pyplot as plt

from notebooks.lib import read_sp

c_sp = read_sp("../data/stick-probability")
by_run = c_sp.groupby(['probability', 'N'])

by_probability = by_run.agg(
    overall_fd=('fd', lambda fd: np.mean(fd[-100:])),
    overall_fd_std=('fd', 'std')
).reset_index().groupby('probability')

ggg = by_probability.agg(
    fd=('overall_fd', 'mean'),
    # TODO Check stats
    fd_std=('overall_fd_std', lambda std: np.sqrt(np.mean(np.square(std))))
)

plt.fill_between(ggg.index, ggg.fd - ggg.fd_std, ggg.fd + ggg.fd_std, alpha=0.2, label=f"Standard error band")
plt.plot(ggg.index, ggg.fd, color='tab:blue', label='fd mean, seeds = 100')
plt.plot(ggg.index, ggg.fd)

plt.xlabel("$p_{stick}$")
plt.ylabel("$fd$")
plt.legend()

plt.savefig('../figures/sp-fd.svg')
plt.savefig('../figures/sp-fd.png')
plt.show()