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_r.txt",
        "particle_1_r.txt",
    ]
    outputs = [
        "../latex/images/simulate_2_particles_no_interaction_xy.pdf",
        "../latex/images/simulate_2_particles_interaction_xy.pdf",
    ]
    labels = [
        r"$p_1$",
        r"$p_2$",
        r"$\hat{p}_1$",
        r"$\hat{p}_2$",
    ]
    colors = [
        "seagreen",
        "darkred"
    ]
    start_pos = set()
    for dir, output in zip(directories, outputs):
        fig, ax = plt.subplots()
        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])
                ax.plot(r[:,0], r[:,1], label=label, color=color)
                start_pos.add((r[0,0],r[0,1]))
        
        for pos in start_pos:
            ax.plot(*pos, "o", color="black")

        ax.set_xlabel(r"x $(\mu m)$")
        ax.set_ylabel(r"y $(\mu m)$")
        ax.legend(loc="upper right")
        ax.axis("equal")
        fig.savefig(output, bbox_inches="tight")


if __name__ == "__main__":
    main()
