import numpy as np
import matplotlib.pyplot as plt
import h5py
from argparse import ArgumentParser
plt.switch_backend("agg")

def F_for_run(run):
    with h5py.File(run, "r") as r:
        F = r["train_F"][...]
    return F

def plot_Fs(Fs, fname, ylim_from_epoch):
    plt.plot(Fs[ylim_from_epoch:], label=fname)

if __name__ == "__main__":
    p = ArgumentParser()
    p.add_argument("runs", nargs="+", help="run files")
    p.add_argument("--ylim-from-epoch", type=int, default=5, help="skip the first E epochs when evaluating plot ylims", metavar="E")
    p.add_argument("--title", default="Free energies", help="skip the first E epochs when evaluating plot ylims", metavar="E")
    args = p.parse_args()
    for Fs, run in zip(map(F_for_run, args.runs), args.runs):
        plot_Fs(Fs, run, args.ylim_from_epoch)
    plt.legend()

    plt.title(args.title)
    plt.savefig("Fs.svg")
