import glob

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


def make_figure(length_index, length, forces, integrators, computes):
    fig, axs = plt.subplots(nrows=len(forces), ncols=len(computes), figsize=(40, 24))
    fig.suptitle(length, fontsize=30)
    for index_compute, compute in enumerate(computes):
        for index_forces, force in enumerate(forces):
            print("Reading file")
            legend = []
            plt.sca(axs[index_forces][index_compute])
            for index_integrator, integrator in enumerate(integrators):
                for filename in glob.glob(f'data/out/*{force}*{integrator}*L{length_index}*_{compute}.dat'):
                    file = pd.read_csv(filename, delimiter=',')
                    file.dropna("columns")
                    ax = plt.plot(file["time"], file["val"])
                    #ax.fill_between(file["time"], y1=file["val"] - file["std"], y2=file["val"] + file["std"], alpha=.5) 
                    plt.ylabel(compute)
                    legend.append(integrator)

                    if index_integrator == 0 and not np.any(file["target"].isnull()):
                        plt.plot(file["time"], file["target"])
                        legend.append("analytisch")

            plt.legend(title='iterators', labels=legend)
    for ax1 in axs:
        for ax in ax1:
            ax.axvline(1.0)
    pad = 5  # in points

    for ax, col in zip(axs[0], computes):
        ax.annotate(col, xy=(0.5, 1), xytext=(0, pad),
                    xycoords='axes fraction', textcoords='offset points',
                    size='large', ha='center', va='baseline', fontsize=25)

    for ax, row in zip(axs[:, 0], forces):
        ax.annotate(row, xy=(0, 0.5), xytext=(-ax.yaxis.labelpad - pad, 0),
                    xycoords=ax.yaxis.label, textcoords='offset points',
                    size='large', ha='right', va='center', fontsize=25)

    fig.tight_layout()
    fig.savefig(f"fig{length_index}.png")
    plt.show()

def main():
    forces = ["harmonic"]
    integrators = ["Euler", "Heun", "Exact", "BDAS", "MBD"]
    computes = ["msd", "oaf", "empxx", "empyy", "x", "x_squared"]
    for l_i, l in enumerate(["sphere"]): #["sphere", "L = 1.5", "L = 2.0"]
        make_figure(length=l, length_index=l_i, forces=forces, computes=computes, integrators=integrators)


if __name__ == "__main__":
    main()