#!/usr/bin/env python

"""
Usage:
  plot_tasks.py TASKTIMERS output.png [Ticks per ms]
  
  where TASKTIMERS is a task timers file generated by quicksched.
  
  This file is part of QuickSched.
  Copyright (c) 2017 Aidan Chalk (aidan.chalk@stfc.ac.uk)
  
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU Lesser General Public License as published
by the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this program.  If not, see <http://www.gnu.org/licenses/>.
"""

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.collections as collections
import matplotlib.cm as cmx
import matplotlib.gridspec as gridspec
matplotlib.use('Agg')
import pylab as pl
import numpy as np
import sys
import math

#  Basic plot configuration.
PLOT_PARAMS = {"axes.labelsize": 10,
               "axes.titlesize": 10,
               "font.size": 14,
               "legend.fontsize": 12,
               "xtick.labelsize": 11,
               "ytick.labelsize": 11,
               "figure.figsize" : (16., 4.),
               "figure.subplot.left" : 0.03,
               "figure.subplot.right" : 0.995,
               "figure.subplot.bottom" : 0.1,
               "figure.subplot.top" : 0.99,
               "figure.subplot.wspace" : 0.,
               "figure.subplot.hspace" : 0.,
               "lines.markersize" : 6,
               "lines.linewidth" : 3.
               }
               
TASKCOLOURS = ["black",
               "lightblue",
               "greenyellow",
               "navy",
               "greenyellow",
               "navy",
               "indigo",
               "cyan",
               "cyan",
               "green",
               "yellow",
               "magenta",
               "mediumorchid",
               #"mediumnightblue",
               "mediumturquoise",
               "mediumvioletred",
               "darkred",
               "darkblue"
               #,"powerblue"
               ]
               
pl.rcParams.update(PLOT_PARAMS)

if len(sys.argv) == 2 and (sys.argv[1][0:2] == "-h" or sys.argv[1][0:3] == "--h"):
    from pydoc import help
    help( "__main__")
    sys.exit(0)

#Input from command line
if len(sys.argv) != 3 and len(sys.argv) != 4:
    print "Usage: ", sys.argv[0], "TASKTIMERS output.png [ticks-per-ms]"
    sys.exit(1)
    

    
infile = sys.argv[1]
outpng = sys.argv[2]
tpms = 0
if len(sys.argv) == 4:
    tpms = int(sys.argv[3])
    
#read input.
data = pl.loadtxt(infile)
nthread=int(max(data[:,1])) + 1

#Recompute timings in MS if possible.
if(tpms != 0):
    data[:,2] = data[:,2] / tpms
    data[:,3] = data[:,3] / tpms
    
start_t = min(data[:,2])
data[:,2] -= start_t
data[:,3] -= start_t
end_t = (max(data[:,3]))
start_t = min(data[:,2])
delta_t = end_t - start_t



fig = pl.figure()
ax = fig.add_subplot(1,1,1)
ax.set_xlim(-delta_t * 0.03, delta_t * 1.03)
ax.set_ylim(0, nthread)

print len(data[:,0])

for j in range(nthread):
    tictocs = np.zeros(len(data[:,0])*2) 
    types = np.zeros(len(data[:,0])*2) 
    for i in range(len(data[:,0])):
        if data[i,1] == j:
            tictocs[i*2+0] = data[i,2]
            tictocs[i*2+1] = data[i,3]
            types[i*2+0] = int(data[i,0])
            types[i*2+1] = int(data[i,0])
        #print tictocs
        #print data[i,1]+0.05
        
    for i in range(len(TASKCOLOURS)):
        print i
        print types==i
        collection = collections.BrokenBarHCollection.span_where(tictocs, ymin=j+0.05, ymax=j+0.95, where=types==i, facecolor=TASKCOLOURS[i], linewidths=1, edgecolor="white")
        ax.add_collection(collection)
    
# Start and end of time-step
ax.plot([start_t, start_t], [0, nthread], 'k--', linewidth=1)
ax.plot([end_t, end_t], [0, nthread], 'k--', linewidth=1)

ax.set_xlabel("time [ms or ticks]")
ax.set_ylabel("Thread ID" )
ax.set_yticks(pl.array(range(nthread)), True)

pl.savefig(outpng, format = 'png', dpi = 1000)
pl.savefig('tasks.eps', format = 'eps', dpi = 1000)

pl.show()

print "Graphics done, output written to", outpng

sys.exit(0)