import numpy as np
from matplotlib import pyplot as plt

from lib.lib import mean_across, read_xyz_alt, read_load

c_direct_neighbours = read_load("../data/c-3d-direct-neighbours", read_xyz_alt)
direct_meaned_by_N = mean_across(c_direct_neighbours)

rust_direct = read_load("../data/rust-3d", read_xyz_alt)
rust_direct_meaned_by_N = mean_across(rust_direct)

rust_offaxis = read_load("../data/rust-3d-offaxis", read_xyz_alt)
rust_offaxis_meaned_by_N = mean_across(rust_direct)

fig, ax = plt.subplots(figsize=(6, 6))


def nc_fd(df, label, color=None, ignore_prefix=0, no_error=False, **kwargs):
    if not no_error:
        plt.fill_between(
            df.N,
            # TODO Check error math here
            (df[ignore_prefix:]['fd']['mean'] - df[ignore_prefix:]['fd']['stderr']),
            (df[ignore_prefix:]['fd']['mean'] + df[ignore_prefix:]['fd']['stderr']),
            color=color,
            alpha=0.2,
            label=f"{label}, standard error band"
        )

    plt.plot(
        df.N,
        df['fd']['mean'],
        color=color,
        label=f"{label}, fd mean",
        **kwargs
    )


# %%
nc_fd(direct_meaned_by_N, label="IPC + 3D Direct", color="tab:blue", no_error=True)
nc_fd(rust_direct_meaned_by_N, label="NF 3D Direct", color="tab:orange", no_error=True)
nc_fd(rust_offaxis_meaned_by_N, label="NF 3D Off-axis", color="tab:purple", linestyle='dashdot', no_error=True)

plt.plot([np.min(c_direct_neighbours.N), np.max(c_direct_neighbours.N)], [2.5, 2.5], color='tab:red', label='Literature')
# plt.fill_between(c_direct_neighbours.N, 2.5 - 0.01, 2.5 + 0.01, alpha=0.2, color='tab:red', label='Literature error band')

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

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