#!/usr/bin/python
# coding=utf8
"""
Compute the cpu/mem/iowait usage rate and draw its line chart.
use: python vm_pic.py [vmstat_ip1.txt|vmstat_ip2.txt|vmstat_ip3.txt...]
     python vm_pic.py
output: vmstat_[ip]_detail.csv\vmstat_[ip].png\summary.csv
Email: [email protected]
"""

import sys
import matplotlib as mpl

mpl.use('Agg')
# import matplotlib.pylab as pl
import matplotlib.pyplot as pl
from pylab import *

mpl.rcParams['font.sans-serif'] = ['SimHei']

print pl.style.available

import numpy as np
import os
import pdb
import time

ltime = []
mem = []
cpu = []
io = []
thresholdDict = {}
ip_over_threshold = {}


def readfile(filename, folder):
    try:
        f1 = open(filename, 'r')
        detailfile = filename[:-4] + '_detail.csv'
        f3 = open(detailfile, 'w')
        f3.write('date,time,mem(%),cpu(%),io(%)\n')
        line = f1.readline()  # 'total mem: '
        # print line
        tmem = float(line.strip().split()[2])
        line = f1.readline()  # 'time procs '
        # print line
        line = f1.readline()  # 'time r b '
        # print line
        line = f1.readline()
        sum_mem = 0
        sum_cpu = 0
        sum_io = 0
        max_mem = 0
        max_cpu = 0
        max_io = 0
        min_mem = 100
        min_cpu = 100
        min_io = 100
        while line:
            # print line
            tmp = line.split()

            tmp_date = tmp[0].strip()  # date

            tmp_time = tmp[1].strip()
            ltime.append(tmp_time)  # time list

            # tmp[5]:free, tmp[6]:buff, tmp[7]:cache, tmp[19]:shm
            fmem = tmem - int(tmp[5].strip()) - int(tmp[6].strip()) - int(tmp[7].strip()) + int(tmp[19].strip())
            tmp_mem = (fmem / tmem) * 100
            mem.append(tmp_mem)  # mem list
            sum_mem += tmp_mem
            if tmp_mem > max_mem:
                max_mem = tmp_mem
            elif tmp_mem < min_mem:
                min_mem = tmp_mem

            tmp_cpu = 100 - int(tmp[16].strip())
            cpu.append(tmp_cpu)  # cpu list
            sum_cpu += tmp_cpu
            if tmp_cpu > max_cpu:
                max_cpu = tmp_cpu
            elif tmp_cpu < min_cpu:
                min_cpu = tmp_cpu

            if tmp_cpu >= int(thresholdDict[folder]):
                if not ip_over_threshold.has_key(folder + '-' + filename[7:-4]):
                    ip_over_threshold[folder + '-' + filename[7:-4]] = True

            tmp_io = int(tmp[17].strip())
            io.append(tmp_io)  # io list
            sum_io += tmp_io
            if tmp_io > max_io:
                max_io = tmp_io
            elif tmp_io < min_io:
                min_io = tmp_io

            f3.write('%s,%s,%f,%d,%d\n' % (tmp_date, tmp_time, tmp_mem, tmp_cpu, tmp_io))
            line = f1.readline()
        f1.close()

        cnt = float(len(ltime))
        avg_mem = sum_mem / cnt
        avg_cpu = sum_cpu / cnt
        avg_io = sum_io / cnt
        f3.write('avg,,%f,%f,%f\n' % (avg_mem, avg_cpu, avg_io))
        f3.write('max,,%f,%f,%f\n' % (max_mem, max_cpu, max_io))
        f3.write('min,,%f,%f,%f\n' % (min_mem, min_cpu, min_io))

        ip = filename[7:-4]
        f4.write('%s,%f,%f,%f,%f,%f,%f,%f,%f,%f\n' % (
            ip, avg_mem, max_mem, min_mem, avg_cpu, max_cpu, min_cpu, avg_io, max_io, min_io))

        f3.close()
    # except:
    #     f2.write('ERROR: %s, The file %s has something wrong, please check it.' % (
    #         time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), filename))
    finally:
        f1.close()
        f3.close()


def pic(filename, folder):
    picname = filename[:-4] + '.png'
    pl.style.use("seaborn-notebook")
    pl.axes([0.1, 0.2, 0.8, 0.7])
    if len(ltime) > 20:
        stp = len(ltime) / 40
    else:
        stp = 1
    xcnt = np.arange(len(ltime), step=stp)
    xlbl = []
    for i in xcnt:
        xlbl.append(ltime[i])
    # pl.xticks(range(len(ltime)),ltime,rotation='vertical')
    pl.xticks(xcnt, xlbl, rotation='vertical')
    pl.yticks(range(0, 101, 20))
    plot1 = pl.plot(cpu, 'g-,', label='cpu', linewidth=1)
    plot2 = pl.plot(mem, 'b', label='mem')
    plot3 = pl.plot(io, 'g', label='i/o wait')
    plot4 = pl.axhline(y=thresholdDict[folder], color='r', label='Threshold', linewidth=0.5)
    pl.title(U'CPU利用率')
    pl.xlabel(U'时间')
    pl.ylabel(U'利用率')
    pl.ylim(0, 100)
    pl.legend()
    pl.grid()
    pl.show()
    pl.savefig(picname)
    pl.close()


if __name__ == '__main__':
    f2 = open('error.txt', 'w')
    f4 = open('summary.csv', 'w')
    f4.write('ip,avg_mem(%),max_mem(%),min_mem(%),avg_cpu(%),max_cpu(%),min_cpu(%),avg_io(%),max_io(%),min_io(%)\n')
    vm_files = []
    if (len(sys.argv) < 2):
        # process all vmstat*.txt in this dir
        cdir = os.getcwd()
        print 'cdir: %s' % (cdir)
        folderlist = os.listdir(cdir)
        for folder in folderlist:
            if os.path.isdir(folder) & os.path.exists(folder + '/.sshvmstat'):
                os.chdir(folder)
                conf_file = open('.sshvmstat', 'r')
                line = conf_file.readline()
                thresholdDict[folder] = line.split(":")[1]
                filelist = os.listdir(cdir + '/' + folder)
                for vmf in filelist:
                    print 'vmf:%s' % (vmf)
                    if vmf[:7] == 'vmstat_' and vmf[-4:] == '.txt':
                        readfile(vmf, folder)
                        pic(vmf, folder)
                        ltime = []
                        cpu = []
                        mem = []
                        io = []
                        print "%s, file %s finished !" % (time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), vmf)
                os.chdir(cdir)
    else:
        for i in range(1, len(sys.argv)):
            vm_files.append(sys.argv[i])
    for vmf in vm_files:
        readfile(vmf)
        pic(vmf)
        ltime = []
        cpu = []
        mem = []
        io = []
        print "%s, file %s finished !" % (time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), vmf)

    f4.write("\nhost over threshold:\n")
    for key in ip_over_threshold:
        f4.write("%s,%s" % (key.split("-")[0], key.split("-")[1]))

    f2.close()
    f4.close()