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

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='plot learnign curve.')
    parser.add_argument('history_file', type=str,
                        help="path to history file.")
    parser.add_argument('--plot_style', nargs='*', default=[],
                        help='plot styles to be used')
    parser.add_argument('--save', default='',
                        help='save the plot in the given file')
    args = parser.parse_args()

    if args.plot_style:
        plt.style.use(args.plot_style)

    df = pd.read_csv(args.history_file)

    # Plot MAE
    fig, ax = plt.subplots()
    ax.plot(df['epoch']+1, df['mae'], label='train', color='blue')
    ax.set_xlabel('epoch')
    ax.set_ylabel('MAE (years)', color='blue')
    ax.set_ylim((8, 14))
    axt = ax.twinx()

    # Plot learning rate
    axt.step(df['epoch']+1, df['lr'], label='train', alpha=0.4, color='k')
    axt.set_yscale('log')
    axt.set_ylabel('learning rate', alpha=0.4, color='k')
    axt.set_ylim((1e-8, 1e-2))

    if args.save:
        plt.savefig(args.save)
    else:
        plt.show()