#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
A small hack to sample process statistics from procfs.

Typically used with the VmRSS and RssAnon metrics.
"""

from pathlib import Path
import sys
import time
import subprocess
from typing import List, Optional

def read_proc(pid: int) -> dict:
    f = Path('/proc') / str(pid) / 'status'
    result = {}
    for line in open(f).readlines():
        k,v = line.split(':')
        result[k.strip()] = v.strip()

    return result


def plot_it(results: "numpy.ndarray",
            out_file: Optional[Path],
            labels: List[str]):
    import matplotlib
    if out_file is not None:
        matplotlib.use('Agg')

    from matplotlib import pyplot as pl
    import numpy as np
    cols = results.shape[1]
    t = results[:,0] - results[0,0]
    for col, label in zip(range(1, cols), labels):
        pl.subplot(cols-1,1,col)
        y = results[:,col]
        pl.scatter(t, y)
        pl.ylim(0, y.max()*1.1)
        pl.ylabel(label)

    pl.xlabel('time (seconds)')
    if out_file:
        pl.savefig(out_file)
    else:
        pl.show()

def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('-m', '--metric', action='append', type=str, help='metric names')
    parser.add_argument('-s', '--period', type=float, help='sampling period (seconds)', default=0.1)
    parser.add_argument('-o', '--output', type=argparse.FileType('w'), help='output file', default=sys.stdout)
    parser.add_argument('-P', '--plot', nargs='?', help='plot it', default=False, const=True)
    parser.add_argument('-p', '--pid', type=int, help='process id')
    parser.add_argument('cmdline', type=str, nargs='*', help='command-line to run')
    args = parser.parse_args()

    metrics = args.metric
    period = args.period
    output = args.output
    if args.pid is not None and args.cmdline is not None:
        raise ValueError("Both pid and command-line given")
    elif args.pid is not None:
        pid = args.pid
    elif args.cmdline == []:
        raise ValueError("Expected either pid or (non-empty) command-line")
    else:
        subp = subprocess.Popen(args.cmdline)
        pid = subp.pid

    if args.plot and output == sys.stdout:
        print("Must output to file in order to plot", file=sys.stderr)
        sys.exit(1)

    try:
        while True:
            t = time.time()
            result = read_proc(pid)
            values = [result[metric].split()[0] for metric in metrics]
            output.write('\t'.join([str(t)] + values) + '\n')
            output.flush()
            time.sleep(period)
    except Exception as e:
        import traceback
        traceback.print_exc(file=sys.stderr)

    if args.plot is not False:
        if args.plot is True:
            out_file = None
        else:
            out_file = args.plot

        import numpy as np
        results = np.genfromtxt(args.output.name)
        plot_it(results, out_file, labels=metrics)

if __name__ == '__main__':
    main()