import os
import json
import matplotlib.pyplot as plt


def tendency_p(log: dict, path: str, higher_is_better=False):
    epochs = [dic['epoch'] for dic in log]
    fig, ax1 = plt.subplots()
    if 'qm7' in path:
        factor = 223.918853
    else:
        factor = 1
    train_p = [dic['train_metric'] * factor for dic in log]
    valid_p = [dic['validate_metric'] * factor for dic in log]
    test_p = [dic['test_metric'] * factor for dic in log]
    ps = zip(valid_p, test_p)
    ps = sorted(ps, key=lambda x: x[0], reverse=higher_is_better)
    print('{}: {:.4f}'.format(path, ps[0][1]))

    ax1.plot(epochs, train_p, color='red', linestyle='--')
    ax1.plot(epochs, test_p, color='red')

    plt.savefig(path)
    plt.close(fig)


tuples = [
    ('qm7', 'NP', False),
]

for d, f, h in tuples:
    if not os.path.exists(d):
        os.mkdir(d)
    json_path = f'{d}/{f}.json'
    graph_path = f'{d}/{f}.png'
    try:
        with open(json_path) as fp:
            log = json.load(fp)
    except FileNotFoundError:
        continue
    tendency_p(log, graph_path, h)
