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/simulate_2_particles/no_interaction/",
        "output/simulate_2_particles/with_interaction/",
    ]
    files = [
        "particle_0",
        "particle_1",
    ]
    labels = [r"$p_1$", r"$p_2$"]
    output = [
        "../latex/images/phase_space_no_interaction",
        "../latex/images/phase_space_interaction",
    ]
    colors = [
        "seagreen",
        "darkred",
    ]

    for i, dir in enumerate(directories):
        fig1, axs1 = plt.subplots(sharex=True)
        fig2, axs2 = plt.subplots(sharex=True)
        for j, (file, label, color) in enumerate(zip(files, labels, colors)):
            r = []
            v = []
            with open(dir + file + "_r.txt") as f:
                lines = f.readlines()
                r = np.array(
                    [list(map(float, line.strip().split(","))) for line in lines]
                )

            with open(dir + file + "_v.txt") as f:
                lines = f.readlines()
                v = np.array(
                    [list(map(float, line.strip().split(","))) for line in lines]
                )
            axs1.plot(r[:, 0], v[:, 0], label=label, color=color)
            axs1.plot(r[0,0], v[0,0], "ko")
            axs2.plot(r[:, 2], v[:, 2], label=label, color=color)
            axs2.plot(r[0,2], v[0,2], "ko")
            axs1.axis("equal")
            axs2.axis("equal")

        axs1.set(xlabel=r"$x (\mu m)$", ylabel=r"$v_x (\mu m / \mu s)$")
        axs2.set(xlabel=r"$z (\mu m)$", ylabel=r"$v_z (\mu m / \mu s)$")

        axs1.legend(loc="upper right")
        axs2.legend(loc="upper right")

        fig1.savefig(f"{output[i]}_x.pdf", bbox_inches="tight")
        fig2.savefig(f"{output[i]}_z.pdf", bbox_inches="tight")
    # plt.show()


if __name__ == "__main__":
    main()
