import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import argparse
import json
import seaborn
seaborn.set_style('whitegrid')

parser = argparse.ArgumentParser(description='Process logs')
parser.add_argument('-x', help='name of parameters on x axis', default='time')
parser.add_argument('-y', help='name of parameters on y axis', default='rmse')
parser.add_argument('f', help='list of logs to parse', nargs='+')
parser.add_argument('-dest', help='name of figure to save', default='fig.png')
args = parser.parse_args()


def default():
    for file_name in args.f:
        x, y = [], []
        with open(file_name) as f:
            for line in f:
                jason = json.loads(line)
                x.append(jason[args.x])
                y.append(jason[args.y])
        plt.plot(x, y, label=file_name)

    plt.legend()
    plt.xlabel(args.x)
    plt.ylabel(args.y)
    plt.savefig(args.dest)

def tmp_finale():
    plt.clf()
    for file_name in args.f:
        x, y = [], []
        with open(file_name) as f:
            for index, line in zip(range(15), f):
                assert index <= 15
                jason = json.loads(line)
                x.append(jason[args.x])
                y.append(np.sqrt(jason[args.y]))
        plt.plot(x, y, label=file_name)

    plt.legend(fontsize=15)
    plt.xlabel(r"Nbr of Batches (size $100$)", fontsize=15)
    plt.ylabel(r"RMSE", fontsize=15)
    plt.savefig(args.dest)

if __name__ == "__main__":
    tmp_finale()