from matplotlib import pyplot as plt

from lib.lib import read_sp, aggregate_sp_fd

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

# %%
c_data = aggregate_sp_fd(c_sp)
rust_data = aggregate_sp_fd(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()