from matplotlib import pyplot as plt

from lib.lib import read_sp, read_xyz_alt, read_xy_alt, aggregate_sp_fd

# %%
d2_raw = read_sp("../data/rust-sticking-probability", read_xy_alt)
d2 = aggregate_sp_fd(d2_raw)
d3_raw = read_sp("../data/rust-3d-offaxis-sp", read_xyz_alt)
d3 = aggregate_sp_fd(d3_raw)

# %%
plt.plot(
    d2.index,
    d2.fd,
    color='tab:blue',
    label='D2'
)

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

plt.plot(
    d3.index,
    d3.fd,
    color='tab:orange',
    label='D3'
)

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

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

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