import matplotlib.pyplot as plt

def main():
    files = [
        "output/time_dependent_potential/f_0.100000.txt",
        "output/time_dependent_potential/f_0.400000.txt",
        "output/time_dependent_potential/f_0.700000.txt",
    ]
    vals = [
        .1,
        .4,
        .7
    ]
    for i in range(3):
        with open(files[i]) as f:
            lines = f.readlines()
            x = []
            y = []
            for line in lines:
                a,b = line.strip().split(",")
                x.append(float(a))
                y.append(float(b))
            plt.plot(x,y,label=f"amplitude: {vals[i]}")

    plt.xlabel(r"$\omega_V$")
    plt.ylabel(r"Fraction of particles left")
    plt.title(r"The fraction of particles left in the Penning trap "
              "after 500 microseconds for different amplitudes and frequencies")
    plt.legend()
    plt.show()

if __name__ == "__main__":
    main()