import argparse
import re
import os
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime

def extract_learning_curves(args):
    paths = args.log_path.split(',')
    if len(paths) == 1 and os.path.isdir(paths[0]):
        paths = [os.path.join(paths[0], f) for f in os.listdir(paths[0]) if os.path.isfile(os.path.join(paths[0], f))]
    learning_curves = {}
    for path in paths:
        learning_curve = []
        lines = open(path, 'r').readlines()
        for line in lines:
            matched = re.match(r'[0-9\- :,]*\[INFO\] - \[Epoch ([0-9]+)\].*Valid MRR: ([0-9\.]+).*', line)
            if matched:
                learning_curve.append(float(matched.group(2)))
                if int(matched.group(1)) >= args.num_epochs:
                    break
        learning_curves[os.path.basename(path)] = learning_curve
    return learning_curves

def draw_learning_curves(args, learning_curves):
    for name in learning_curves.keys():
        epochs = np.arange(len(learning_curves[name]))
        matched = re.match(r'(.*)\..*', name)
        if matched:
            label = matched.group(1)
        else:
            label = name
        plt.plot(epochs, learning_curves[name], label = label)
    plt.xlabel("Epochs")
    plt.ylabel("MRR")
    plt.legend(title=args.legend_title)
    plt.savefig(os.path.join(args.out_path, str(round(datetime.utcnow().timestamp() * 1000)) + '.' + args.fig_filetype))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Parser For Arguments", formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--task', default = None, type=str)
    parser.add_argument('--log_path', type=str, default=None)
    parser.add_argument('--out_path', type=str, default=None)
    parser.add_argument('--num_epochs', type=int, default=200)
    parser.add_argument('--legend_title', type=str, default="Learning rate")
    parser.add_argument('--fig_filetype', type=str, default="svg")

    args = parser.parse_args()

    if (args.task == 'learning_curve'):
        draw_learning_curves(args, extract_learning_curves(args))