import matplotlib.pyplot as plt
import numpy as np
import pandas as pd


def main():
    time = np.arange(0.0, 4.0, 0.01)
    xstart = 1.0
    ystart = 1.0
    x_mean = xstart * np.exp(-time)
    y_mean = ystart * np.exp(-time)
    x_sq_mean = xstart * xstart * np.exp(-2 * time) + (1 - np.exp(-2 * time))
    y_sq_mean = ystart * ystart * np.exp(-2 * time) + (1 - np.exp(-2 * time))

    x = pd.read_csv("./data/out/harmonic_force_euler_L0_x.dat")
    x_sq = pd.read_csv("./data/out/harmonic_force_euler_L0_x_squared.dat")
    msd = pd.read_csv("./data/out/harmonic_force_euler_L0_msd.dat")

    msd_mean = (1 - np.exp(-2 * time))
    msd_2_mean = 2 * (1 - 2 * np.exp(- time) + np.exp(-2 * time)) + 2 * (1 - np.exp(-2 * time))

    plt.plot(time, x_mean, label="a <x>")
    plt.plot(x["time"], x["val"], label="<x>")

    plt.plot(time, x_sq_mean, label="a <x²>")
    plt.plot(x_sq["time"], x_sq["val"], label="<x²>")

    plt.plot(time, 2 * msd_mean, label="a msd")
    plt.plot(time, msd_2_mean, label="a msd")

    plt.plot(msd["time"], msd["val"], label="msd")

    # plt.plot(x["time"], x_sq["val"] - np.power(x["val"], 2.0))
    plt.plot(x["time"], 2 * (x_sq["val"] - np.power(x["val"], 2.0)), label="<x²>-<x>²")
    plt.legend(loc=1)
    plt.show()


if __name__ == "__main__":
    main()