import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns

sns.set_theme()
params = {
    "font.family": "Serif",
    "font.serif": "Roman", 
    "text.usetex": True,
    "axes.titlesize": "large",
    "axes.labelsize": "large",
    "xtick.labelsize": "large",
    "ytick.labelsize": "large",
    "legend.fontsize": "medium"
}
plt.rcParams.update(params)

def main():
    directories = [
        "output/relative_error/RK4/",
        "output/relative_error/euler/",
    ]
    files = [
        "4000_steps.txt",
        "8000_steps.txt",
        "16000_steps.txt",
        "32000_steps.txt",
    ]
    labels = [
        r"$n_1$",
        r"$n_2$",
        r"$n_3$",
        r"$n_4$",
    ]
    titles = [
        "Relative error for the RK4 method",
        "Relative error for the forward Euler method"
    ]
    methods = [
        "rk4",
        "euler"
    ]
    fig1, axs1 = plt.subplots(2,1)
    for i, (dir, title) in enumerate(list(zip(directories, titles))):
        max_err = []
        for label, file in zip(labels, files):
            with open(dir+file) as f:
                lines = f.readlines()
                t = np.linspace(0, 50, len(lines))
                r = np.array([float(line.strip()) for line in lines])
                max_err.append(max(r))
                axs1[i].plot(t, r, label=label)
        
        axs1[i].set(xlabel=r"t $(\mu s)$", ylabel = r"relative_error $(\mu m)$")
        axs1[i].legend()
        axs1[i].set_title(title)

        conv_rate = 1/3 * sum([np.log2(max_err[i+1]/max_err[i])/np.log2(.5) for i in range(3)])
        print(f"{methods[i]}: {conv_rate}")

    fig1.savefig("../latex/images/phase_space_2_particles_x.pdf")


if __name__ == "__main__":
    main()