import matplotlib.pyplot as plt
import numpy as np

def main():
    directories = {
        "output/simulate_2_particles/no_interaction/",
        "output/simulate_2_particles/with_interaction/",
    }
    titles = {
        "particles without interaction",
        "particles with interaction"
    }
    files = [
        "particle_0_r.txt",
        "particle_0_v.txt",
        "particle_1_r.txt",
        "particle_1_v.txt",
    ]
    labels = [
        r"particle 1 r",
        r"particle 1 v",
        r"particle 2 r",
        r"particle 2 v",
    ]
    colors = [
        "lightskyblue",
        "deepskyblue",
        "salmon",
        "tomato",
    ]
    fig1, axs1 = plt.subplots(2,1)
    fig2, axs2 = plt.subplots(2,1)
    for i, (dir, title) in enumerate(zip(directories, titles)):
        for label, color, file in zip(labels, colors, files):
            with open(dir+file) as f:
                lines = f.readlines()
                t = np.linspace(0, 50, len(lines))
                r = np.array([list(map(float, line.strip().split(","))) for line in lines])
                axs1[i].plot(t, r[:,0], label=label, color=color)
                axs2[i].plot(t, r[:,2], label=label, color=color)
        
            axs1[i].set(xlabel=r"t $(\mu s)$", ylabel = r"z $(\mu m)$")
        
        axs1[i].legend()
        axs1[i].set_title(title)

    # plt.show()
    fig1.savefig("../latex/images/phase_space_2_particles_x.pdf")
    fig2.savefig("../latex/images/phase_space_2_particles_z.pdf")


if __name__ == "__main__":
    main()