import numpy as np
import matplotlib.pyplot as plt
from lib.lib import read_load

alpha = read_load("../data/alpha")
meaned_by_N = alpha.groupby('N').agg({'fd': ['mean', 'std']}) \
    .reset_index() \
    .replace([np.inf, -np.inf], np.nan)

without_prefix = meaned_by_N[50:]
fig, ax = plt.subplots(figsize=(6, 6))

plt.fill_between(
    without_prefix.N,
    # TODO Check error math here
    (without_prefix['fd']['mean'] - without_prefix['fd']['std'] / np.sqrt(20)),
    (without_prefix['fd']['mean'] + without_prefix['fd']['std'] / np.sqrt(20)),
    alpha=0.2, label=f"Standard error band"
)

plt.plot(
    without_prefix.N,
    without_prefix['fd']['mean'],
    color='tab:blue', label='fd mean, seeds = 20'
)

plt.plot([50, 10000], [1.71, 1.71], color='red', label='Theory')
plt.fill_between(without_prefix.N, 1.71 - 0.01, 1.71 + 0.01, alpha=0.2, label='Theory error band')

plt.xlabel("$N_C$")
plt.ylabel("$fd$ (instantaneous)")
plt.legend()

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