import matplotlib.pyplot as plt
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():
    colors = [
        "mediumseagreen",
        "darkred",
        "darkgoldenrod"
    ]
    files = [
        "output/time_dependent_potential/wide_sweep.txt",
        "output/time_dependent_potential/narrow_sweep.txt",
        "output/time_dependent_potential/narrow_sweep_fine.txt",
        "output/time_dependent_potential/narrow_sweep_interactions.txt",
        "output/time_dependent_potential/narrow_sweep_interactions_fine.txt",
    ]
    outputs = [
        "../latex/images/particles_left_wide_sweep.pdf",
        "../latex/images/particles_left_narrow_sweep.pdf",
        "../latex/images/particles_left_narrow_sweep_fine.pdf",
        "../latex/images/particles_left_narrow_sweep_interactions.pdf",
        "../latex/images/particles_left_narrow_sweep_interactions_fine.pdf",
    ]
    for file, output in zip(files, outputs):
        with open(file) as f:
            lines = f.readlines()
            x = []
            y1 = []
            y2 = []
            y3 = []
            for line in lines:
                l = line.strip().split(",")
                x.append(float(l[0]))
                y1.append(float(l[1]))
                y2.append(float(l[2]))
                y3.append(float(l[3]))

            fig, ax = plt.subplots()
            ax.plot(x, y1, label=r"$f_{1} = 0.1$", color=colors[0])
            ax.plot(x, y2, label=r"$f_{2} = 0.4$", color=colors[1])
            ax.plot(x, y3, label=r"$f_{3} = 0.7$", color=colors[2])

        ax.set_xlabel(r"Angular frequency $\omega_V$ (MHz)")
        # ax.set_xlim((0, 2.8))
        ax.set_ylabel(r"Fraction of particles left")
        # ax.set_ylim((-0.1, 1.1))
        # plt.title(r"The fraction of particles left in the Penning trap "
        #           "after 500 microseconds for different amplitudes and frequencies")
        
        ax.legend(loc="upper right")

        # plt.show()
        fig.savefig(output, bbox_inches="tight")

if __name__ == "__main__":
    main()