#!/usr/bin/env python3
import argparse
import logging

try:
    import matplotlib  # noqa
    matplotlib.use('agg')  # noqa
    import pandas as pd
    import matplotlib.pyplot as plt
    import numpy as np
except ImportError:
    import sys
    print("Error: Install pandas and matplotlib to plot!", file=sys.stderr)


def cpuPlot(df, tag):
    fig, ax = plt.subplots()

    # cpu_usage = (df['cpu_time_system'] + df['cpu_time_user'])
    # cpu_usage = cpu_usage - cpu_usage.iloc[0]
    # cpu_usage.rolling('60s').apply(lambda x: x.iloc[-1] - x.iloc[0])
    # iowait_usage = df['iowait'].rolling(
    #     '60s').apply(lambda x: x.iloc[-1] - x.iloc[0])

    cpu_usage = (df['cpu_time_system'] + df['cpu_time_user']
                 ).rolling('1200s').apply(lambda x: np.abs(x.iloc[-1] - x.iloc[0]))

    cpu_usage.plot(label="Average usage (1m)")
    # iowait_usage.plot(label="Average usage (2m)")
    plt.title(f"cpu percentage {tag}")
    plt.ylabel("cpu percentage [%]")
    plt.xlabel("time")
    # lgn = plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

    if tag is None:
        plt.savefig("cpu.png")
    else:
        plt.savefig(f"{tag}_cpu.png",
                    # bbox_extra_artists=[lgn],
                    bbox_inches='tight')
    plt.clf()


def memPlot(df, tag):
    mem_total = df['mem_rss'] + df['mem_vms']
    df['mem'] = (mem_total - mem_total.iloc[0]) * 1e-9
    df['mem'].rolling('60s').mean().plot(
        label="Average usage (1m)")
    plt.title(f"memory usage {tag}")
    plt.ylabel("memory usage [GB]")
    plt.xlabel("time")
    lgn = plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left')

    if tag is None:
        print("saving mem.png")
        plt.savefig("mem.png")
    else:
        print(f"saving {tag}_mem.png")
        plt.savefig(f"{tag}_mem.png",
                    bbox_extra_artists=[lgn],
                    bbox_inches='tight')

    plt.clf()


def plotter(infile: str, tag: str):
    """
    Plot the results from the runner step. 
    It can be done after the fact with `-j` option to just plot.

    Args:
        infile (str, optional): Input filename for csv. Defaults to "stats.csv".
    """

    # Default to input name without extention
    if tag is None:
        tag = infile[:-4]

    # Load in dataframe and put 0th column as DateTimeIndex
    df = pd.read_csv(infile,
                     index_col=[0],
                     parse_dates=[0])

    cpuPlot(df, tag)
    # memPlot(df, tag)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-t", "--tag", type=str,
                        help="Tags the process and gives a name to the plots and statistcs csv file.",
                        default=None)
    parser.add_argument("-d", "--debug", action="store_true",
                        help="Run with debugging info.",
                        default=False)
    parser.add_argument("-i", "--input", type=str,
                        help="Input statistics file from pagurus",
                        default=None)

    args = parser.parse_args()

    if args.debug:
        logging.basicConfig(
            format='%(asctime)s %(levelname)s ==> %(message)s', level=logging.INFO)
    else:
        logging.basicConfig(level=logging.FATAL)

    logging.info(f'Using tag {args.tag}')

    plotter(infile=args.input,
            tag=f"{args.tag}" if args.tag is not None else None)
