import matplotlib.pyplot as plt
import pandas as pd
import argparse


def line_format(label):
    """
    Convert time label to the format of pandas line plot
    """
    month = label.month_name()[:3]
    if month == 'Jan':
        month = f'{label.year}\n{month}'
    return month


def create_figure(args):
    df = pd.read_csv(args.csv_file, index_col='Account', parse_dates=True)
    df = df.drop(columns="Total:")
    figsize=(10, 6)
    if args.type == 'bar':
        ax = df.plot(grid=True, kind='bar', figsize=figsize)
        ax.set_xticklabels(map(lambda x: line_format(x), df.index))
    else:
        ax = df.plot(grid=True, figsize=figsize)
    plt.savefig(args.png_file)
    # plt.show()


def get_args():
    parser = argparse.ArgumentParser(description='Transform a CSV file into a bar chart.')
    parser.add_argument('csv_file', type=str, help='CSV data file')
    parser.add_argument('png_file', type=str, help='Chart png file')
    parser.add_argument('--type', default='', type=str, help='Chart png file')
    return parser.parse_args()


if __name__ == "__main__":
    args = get_args()
    create_figure(args)