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 z(t):
    V_0 = 25.*9.64852558 * 10**4
    m = 40.
    d = 500.
    w_z = np.sqrt((2.*V_0)/(m*d*d))
    return 20.*np.cos(w_z*t)

def main():
    colors = [
        "lightskyblue",
        "deepskyblue",
        "salmon",
        "tomato",
        "mediumaquamarine",
        "mediumseagreen"
    ]
    filename = "output/simulate_single_particle/particle_0_r.txt"
    r = t = []
    with open(filename) 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])

    fig, ax = plt.subplots()
    ax.plot(t, r[:, 2], label="approximation", color="mediumseagreen")
    ax.plot(t, z(t), label="analytical", color="black", linestyle="dotted")
    ax.set_xlabel(r"t $(\mu s)$")
    ax.set_xlim((-5, 55))
    ax.set_ylabel(r"z $(\mu m)$")
    ax.set_ylim((-25, 25))
    # plt.title(r"Movement of a single particle in the x direction")
    ax.legend(loc="upper right")
    fig.savefig("../latex/images/single_particle.pdf", bbox_inches="tight")
    # plt.show()


if __name__ == "__main__":
    main()