from datetime import datetime, timedelta
from matplotlib.collections import LineCollection
from numpy import array
from pylab import subplot, figure, show

gpuFileNames = ['GPU.txt']
cpuFileNames = ['CPU.txt']
fpsFileNames = ['3DMark06 2010-05-01 12-26-51-93 fps.csv',
                '3DMark06 2010-05-01 12-28-47-27 fps.csv',
                '3DMark06 2010-05-01 12-32-00-08 fps.csv',
                '3DMark06 2010-05-01 12-33-33-37 fps.csv']

# Translation of GPU-Z's field names
gpuFields = {
    u'GPU Core Clock [MHz]'    : ('Core', 'MHz'),
    u'GPU Memory Clock [MHz]'  : ('Memory', 'MHz'),
    u'GPU Temperature [\xb0C]' : ('?', u'\xb0C'),
    u'Fan Speed [%]'           : ('Fan Speed', '%'),
    u'GPU Load [%]'            : ('Load', '%'),
    u'Fan Speed [RPM]'         : ('Fan Speed', 'RPM'),
    u'GPU Temp. #1 [\xb0C]'    : ('DispIO', u'\xb0C'),
    u'GPU Temp. #2 [\xb0C]'    : ('MemIO', u'\xb0C'),
    u'GPU Temp. #3 [\xb0C]'    : ('Shader', u'\xb0C')
}
def gpuHeads(heads, count):
    return [gpuFields[unicode(head, 'latin')] for head in heads]

# Generates CPU headings
def cpuHeads(heads, count):
    heads = []
    heads.append(('', u'MHz'))
    for i in xrange(count - 4):
        heads.append(('Core #{0}'.format(i), u'\xb0C'))
    heads.append(('Load', u'%'))
    return heads

# Generates FPS headings
def fpsHeads(heads, count):
    return [('', u'FPS')]

colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k']

class DataSet(object):
    """ Data set - corresponds to a single set of data """
    def __init__(self, heads, data):
        self.heads = heads
        self.data = array(data)

    def __len__(self):
        return len(self.data)

class DataLog(object):
    """ Data log - corresponds to a single data file """
    def __init__(self, name, dataSets, start, end):
        self.name = name
        self.dataSets = dataSets
        self.start = start
        self.end = end
        self.len = 0
        for dataSet in dataSets:
            self.len += len(dataSet)

    def __len__(self):
        return self.len

def main():
    """ Main function. Loads and plots the data """
    global gpu, cpu, fps
    gpu = loadData(gpuFileNames, 'GPU', ',', 'Date', gpuHeads, 1, '%Y-%m-%d %H:%M:%S')
    cpu = loadData(cpuFileNames, 'CPU', None, 'DATE', cpuHeads, 2, '%m/%d/%y %H:%M:%S')
    fps = loadData(fpsFileNames, 'Frame rate', None, 'FPS', fpsHeads, 0, '%Y-%m-%d %H-%M-%S')

    # Draw the graph
    draw(gpu, cpu, fps)

def loadData(fileNames, logName, sep, firstHead, headsFun, dateLen, dateFormat):
    """ Loads data sets from a file into a DataLog """
    dataSets = []
    data = []
    firstDate = None
    prev = None

    try:
        for fileName in fileNames:
            date = None
            f = open(fileName)

            for line in f:

                toks = line.split(sep)
                if not toks: continue # Empty line
                toks = [tok.strip() for tok in toks] # Strip empty space
                if toks[-1] == '': toks = toks[:-1]

                # Read and interpret data headings
                if toks[0] == firstHead:
                    if data:
                        dataSets.append(DataSet(heads, data))
                    heads, data = None, []
                    headToks = toks[dateLen:]
                    continue
                if heads is None:
                    heads = headsFun(headToks, len(toks))

                # Read or calculate date/time
                if dateLen > 0:
                    date = datetime.strptime(' '.join(toks[0:dateLen]), dateFormat)
                elif date is None:
                    dstr = ' '.join(fileName.rsplit(None, 3)[1:-1]).rsplit('-', 1)[0]
                    date = datetime.strptime(dstr, dateFormat)
                else:
                    date += timedelta(0,1,0)
                if firstDate is None: firstDate = date
                time = (date - firstDate).seconds
                if prev is None: prev = time - 1

                # Leave gaps where appropriate
                if time - prev > 1 and data:
                    dataSets.append(DataSet(heads, data))
                    data = []

                # Read sensor values
                for i in xrange(dateLen, len(toks)):
                    toks[i] = float(toks[i])

                # Catch anomalous data
                valid = True
                for i, v in enumerate(toks[dateLen:]):
                    if v < 0 or v > 10000:
                        print "Anomalous reading for {0}[{1}]: {2} at {3}".format(
                            heads[i][0], heads[i][1], v, date)
                        valid = False
                if valid:
                    data.append([time] + toks[dateLen:])
                    prev = time

            if data:
                dataSets.append(DataSet(heads, data))
                data = []

            f.close()

    except IOError:
        print fileName + " could not be loaded"

    return DataLog(logName, dataSets, firstDate, date)

def draw(gpu, cpu, fps):
    """ Plotting function.
    Takes three DataLogs, for the GPU, CPU and FPS data."""

    # Summarize time period of data
    ss = "{1} lines of {0} data from {2} to {3}"
    start, end = None, None
    for log in [gpu, cpu, fps]:
        if log:
            print ss.format(log.name, len(log), log.start, log.end)
            if not start or start > log.start: start = log.start
            if not end or end < log.end: end = log.end
    for log in [gpu, cpu, fps]:
        offset(start, log)

    # Maximum value, leaving some room for the legend
    xmax = (end-start).seconds * 1.15

    print
    print "{0:17}: Min/Avg/Max".format('Measurement')

    tempFilt = lambda x: x[0] != '?' and x[1] == u'\xb0C'
    clockFilt = lambda x: x[1] == u'MHz'
    loadFilt = lambda x: 'Fan' not in x[0] and x[1] == u'%'
    fpsFilt = lambda x: x[1] == u'FPS'
    graphConds = [([cpu], tempFilt, u'Temperature [\xb0C]'),
                  ([gpu], tempFilt, u'Temperature [\xb0C]'),
                  ([cpu, gpu], clockFilt, u'Clock speed [MHz]'),
                  ([cpu, gpu], loadFilt, u'Load [%]'),
                  ([fps], fpsFilt, u'Frame rate [FPS]')
                  ]
    graphInfo = []
    for cond in graphConds:
        ret = makeLines(cond[0], cond[1])
        if ret:
            graphInfo.append(ret + (cond[2],))
    np = len(graphInfo)

    fig = figure()
    for i in xrange(np):
        lcs, mn, mx, ylabel = graphInfo[i]
        ax = subplot(np, 1, i+1)
        for lc in lcs:
            ax.add_collection(lc)
        ax.grid(True)
        if i == 0:
            ax.set_title('Status from {0} to {1}'.format(start, end))
        ax.set_ylabel(ylabel)
        if i == np-1:
            ax.set_xlabel('Time [s]')
        ax.axis([0, xmax, mn, mx])
        ax.legend()

    show()

def makeLines(dataLogs, filt):
    """ Create line collections from the data logs """
    mn, mx = None, None
    uids, liness = [], []
    hasData = False

    for log in dataLogs:
        for dataSet in log.dataSets:
            curHeads = filter(filt, dataSet.heads)
            if curHeads: hasData = True

            for head in curHeads:
                uid = (log.name, head)
                if uid not in uids:
                    uids.append(uid)
                    liness.append([])
                i = uids.index(uid)
                j = dataSet.heads.index(head)
                liness[i].append(dataSet.data[:, (0, j+1)])

    if not hasData:
        return None

    for i in xrange(len(uids)):
        uid = uids[i]
        cmn, cavg, cmx = doStats(uid, liness[i])
        if mn is None or cmn < mn:
            mn = cmn
        if mx is None or cmx > mx:
            mx = cmx
        liness[i] = LineCollection(liness[i])
        liness[i].set_label(uid[0] + " " + uid[1][0])
        liness[i].set_color(colors[i])

    return liness, mn, mx

def offset(start, dataLog):
    """ Change start date of the log """
    if dataLog:
        delta = (dataLog.start - start).seconds
        for dataSet in dataLog.dataSets:
            dataSet.data[:, 0] += delta
        dataLog.start = start

def doStats(uid, lines):
    """ Return and print min/avg/max stats"""
    mn, mx = None, None
    count, tot = 0, 0

    for line in lines:
        vals = line[:, 1]
        cmn = vals.min()
        if mn is None or cmn < mn:
            mn = cmn
        cmx = vals.max()
        if mx is None or cmx > mx:
            mx = cmx
        tot += vals.sum()
        count += len(vals)
    if count > 0:
        avg = float(tot) / count
    else:
        avg = 0

    title = uid[0]
    if uid[1][0]:
        if title:
            title += " "
        title += uid[1][0]
    units = "[" + uid[1][1] + "]"
    print u"{0:11}{1:>6}: {2:.1f}/{3:.1f}/{4:.1f}".format(title, units, mn, avg, mx)
    return mn, avg, mx

if __name__ == '__main__':
    main()