import numpy as np
from matplotlib import pyplot as plt

from notebooks.lib import read_sp


def process_sp(sp):
    by_run = 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, do we need to do /sqrt(n) here?
        fd_std=('overall_fd_std', lambda std: np.sqrt(np.mean(np.square(std))))
    )

    return ggg


c_sp = read_sp("../data/stick-probability")
rust_sp = read_sp("../data/rust-sticking-probability")

# %%
c_data = process_sp(c_sp)
rust_data = process_sp(rust_sp)

# %%

# plt.fill_between(
#     c_data.index,
#     c_data.fd - c_data.fd_std,
#     c_data.fd + c_data.fd_std,
#     alpha=0.2,
#     color='tab:blue',
#     label=f"IPC + PS, Standard error band"
# )

plt.plot(
    c_data.index,
    c_data.fd,
    color='tab:blue',
    label='IPC + PS, fd mean, seeds = 100'
)

# plt.fill_between(
#     rust_data.index,
#     rust_data.fd - rust_data.fd_std,
#     rust_data.fd + rust_data.fd_std,
#     alpha=0.2,
#     color='tab:orange',
#     label=f"NF, Standard error band"
# )

plt.plot(
    rust_data.index,
    rust_data.fd,
    color='tab:orange',
    label='NF, fd mean, seeds = 100'
)

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

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