"""
Plot training history from careless output.
"""

import seaborn as sns
import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from argparse import ArgumentParser,RawTextHelpFormatter


class ArgumentParser(ArgumentParser):
    def __init__(self, *args, **kwargs):
        super().__init__(
            formatter_class=RawTextHelpFormatter, 
            description=__doc__
        )
        self.add_argument('history_csv', type=str, help='A _history.csv file generated by careless')
        self.add_argument('-o', type=str, default=None, help='Output image file name')
        self.add_argument('-s', '--show', action='store_true')


def plot_history(history_csv, step_key='step'):
    df = pd.read_csv(history_csv).melt(step_key)
    val = df.variable.str.endswith('_val')
    df['variable'] = df.variable.str.removesuffix('_val')
    df['Set'] = np.array(['Train', 'Test'])[val.to_numpy('int')]
    sns.lineplot(
        data=df,
        x='step',
        y='value',
        hue='variable',
        style='Set',
        palette='Dark2',
    )
    plt.semilogy()
    plt.grid(which='both', axis='both', linestyle='--')


def run_analysis(parser):
    plot_history(parser.history_csv)

    if parser.o is not None: 
        plt.savefig(parser.o, dpi=300)
    if parser.show:
        plt.show()

def main():
    parser = ArgumentParser()
    parser = parser.parse_args()
    run_analysis(parser)

