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

#
#  Copyright: 2006 by Dalibo <dalibo.com>
#  Copyright: 2007 Dimitri Fontaine <dim@tapoueh.org>
#  Created: 2006 by Dimitri Fontaine <dim@tapoueh.org>
#
#  Modified: 2008 by Nicolas Niclausse
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 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 General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307, USA.
#
#  In addition, as a special exception, you have the permission to
#  link the code of this program with any library released under
#  the EPL license and distribute linked combinations including
#  the two.

# A plotter for tsung text generated data

import os, sys

LIBDIR         = "/usr/local/lib/tsung"
SHAREDIR       = "/usr/local/share/tsung"
USERDIR        = ".tsung"
CONF           = "http.plots.en.conf"
SYS_STATS_CONF = os.path.join(SHAREDIR, 'tsung_plotter', 'stats.conf')

sys.path.append(LIBDIR)

from tsung_plotter.tsung import TsungLog
from ConfigParser import ConfigParser

# Prevents pylab from requiring a X Server to run
import matplotlib
matplotlib.use("Agg")

from pylab import *


class Plot:
    """ The ploting class, using matplotlib """

    def __init__(self,
                 name    = "",
                 xlabel  = "",
                 ylabel  = "",
                 title   = "",
                 legends = None,
                 position = 0,
                 styles  = ['b-', 'r-', 'g-', 'c-', 'y^', 'kv'],
                 colors  = ['b', 'r', 'g', 'c', 'y', 'k'],
                 dpi     = 150,
                 tn_dpi  = 50,
                 xfactor = 1,
                 yfactor = 1,
                 yscale  = 'linear',
                 plottype  = 'lineplot', # can be lineplot or bar
                 outdir  = '/tmp/tsung-plotter',
                 imgtype = 'png'):

        # Only used if no plot type given
        self.name      = name
        self.xlabel    = xlabel
        self.ylabel    = ylabel
        self.title     = title
        self.legends   = legends
        self.position  = position
        self.dpi       = dpi
        self.thumb_dpi = tn_dpi
        self.xfactor   = xfactor
        self.yfactor   = yfactor
        self.yscale    = yscale
        self.plottype  = plottype
        self.styles    = styles
        self.colors    = colors
        self.outdir    = outdir
        self.imgtype   = imgtype

    def plot(self, stats, dataset):
        """ draw a simple timeline or bar plots from two dataset """

        if self.plottype == "lineplot":
            # prepare matplotlib graph
            clf()
            ax = subplot(111)
            ax.set_yscale(self.yscale)

            count = 0
            for data in dataset:
                # get give type data for plotting
                nstats = 0
                for (name, stat) in stats:
                    pdata = data.stat(name, stat)
                    # we want timestamp sorted data to plot
                    ts = pdata.keys()
                    ts.sort()
                    values = [pdata[t] / self.yfactor[nstats%len(self.yfactor)] for t in ts]

                    if self.xfactor not in [1, "1"]:
                        ts = [x / self.xfactor for x in ts]

                    # Now use matplotlib to do the plotting
                    plot(ts, values, self.styles[count%len(self.styles)])
                    count += 1
                    nstats += 1
            # Now setup the legend
            title(self.title)
            xlabel(self.xlabel)
            ylabel(self.ylabel)
            legend(self.legends, loc=self.position)

            # we want to draw a grid
            grid(True)

            filename = os.path.join(self.outdir,
                                    "%s.%s" % (self.name, self.imgtype))
            thumbnail = os.path.join(self.outdir,
                                     "%s_tn.%s" % (self.name, self.imgtype))

            savefig(filename, dpi=self.dpi, orientation='portrait')
            savefig(thumbnail, dpi=self.thumb_dpi, orientation='portrait')
            return filename
        else:
            # box chart for gmean
            clf()
            width = 1.0/len(dataset)
            nbox=len(dataset)
            ind=arange(1)
            count = 0
            current = {}
            percent = []
            ticksp = []
            ax = subplot(111)
            ax.set_yscale(self.yscale)
            for data in dataset:
                nstats=0
                for (name, stat) in stats:
                    pdata = data.stat(name, stat)
                    size= len(pdata)
                    if size > 0 :
                        # the data we use is the latest value:
                        current[count] = pdata.values()[size-1] / self.yfactor[nstats]
                        ax.bar(ind+count*width, ( current[count] ), width,
                               color=self.colors[count] )
                        # get percent of increase/decrease:
                        if count > 0:
                            percent.append(str(round(-100 + current[count] / current[0] * 100)) + "% ")
                        else:
                            percent.append("reference")
                        currenttick=(0.5 + count)/nbox
                        ticksp.append(currenttick)
                        count += 1
                        nstats += 1

            boxchart = os.path.join(self.outdir,
                                    "%s.%s" % (self.name, self.imgtype))
            boxchart_thumb = os.path.join(self.outdir,
                                    "%s_tn.%s" % (self.name, self.imgtype))
            title(self.title)
            legend(self.legends, loc=self.position)
            ax.set_xticks(ticksp)
            ax.set_xticklabels(percent)
            ylabel(self.ylabel)
            savefig(boxchart, dpi=self.dpi, orientation='portrait')
            savefig(boxchart_thumb, dpi=self.thumb_dpi, orientation='portrait')
            return boxchart


def main(conffile, logs, legends, outdir, verbose):
    """ Produce plots from given """

    dataset = [d for f, d in logs]

    config = ConfigParser()
    config.read(conffile)

    for s in config.sections():
        p = Plot(name = s, outdir = outdir)

        # defaults
        for d, v in config.defaults().items():
            if d != 'encoding':
                p.__dict__[d] = v
            else:
                encoding = v

        # stats
        # conf:   200.count, 400.count
        # result: [["200", "count"], ["400", "count"]]
        try:
            stats = [x.strip().rsplit('.', 1)
                     for x in config.get(s, 'stats').split(' ')]
        except:
            print 'error: unable to read plot "%s" stats' % s
            continue

        if config.has_option(s, 'styles'):
            p.styles = [x.strip() for x in config.get(s, 'styles').split(' ')]

        if config.has_option(s, 'legend'):
            # this is the legend prefix$
            l = []
            count    = 0
            clegends = config.get(s, 'legend').decode(encoding).split(',')

            for f in logs:
                l     += ['%s %s' % (x.strip(), legends[count]) \
                          for x in clegends]
                count += 1
        p.legends = l
        
        if config.has_option(s, 'position'):
            # legend position according to matplotlib standard values (1-10)
            val = config.get(s, 'position').decode(encoding)
            if val.isdigit():
                p.position = int(val)
            else:
                p.position = val
            

        if config.has_option(s, 'yfactor'):
            try:
                p.__dict__['yfactor'] = map(float,config.get(s, 'yfactor').decode(encoding).split(','))
            except ValueError:
                print 'warning: %s yfactor not a number: %s' \
                    % (p.name, config.get(s, yfactor))
        # Text parameters - to decode into specified encoding
        for attr in ['title', 'xlabel', 'ylabel', 'plottype', 'yscale']:
            if config.has_option(s, attr):
                cstring = config.get(s, attr).decode(encoding)
                p.__dict__[attr] = cstring

        # Numerical parameters
        for attr in ['xfactor', 'dpi', 'tn_dpi']:
            if config.has_option(s, attr):
                try:
                    p.__dict__[attr] = config.getfloat(s, attr)
                except ValueError:
                    print 'warning: %s %s not a number: %s' \
                          % (p.name, attr, config.get(s, attr))

        outfile = p.plot(stats, dataset)

        if verbose:
            print 'Generated plot %s' % outfile

if __name__ == "__main__":
    from optparse import OptionParser

    parser = OptionParser()
    parser.add_option("-c", "--config", dest="config", default=None,
                      help="configuration file")
    parser.add_option("-d", "--outdir", dest="outdir", default="/tmp/tsung",
                      help="output dir where to save plots (/tmp/tsung)")
    parser.add_option("-v", action="store_true", dest="verbose", default=False,
                      help="be verbose")

    (options, args) = parser.parse_args()

    if options.config is None:
        userconf = os.path.join(os.environ['HOME'], USERDIR, CONF)
        if os.access(userconf, os.R_OK):
            config = userconf
        else:
            config = os.path.join(SHAREDIR, 'tsung_plotter', CONF)
    else:
        config = options.config

    if options.verbose:
        print 'Using %s configuration file' % config

    if not os.access(config, os.R_OK):
        print "can't read configuration file: %s" % config
        sys.exit(1)

    # FIXME: error control
    # OSError: [Errno 17] Le fichier existe.: '/tmp/tsung'
    try:
        os.makedirs(options.outdir)
    except:
        pass

    # args are legend then file, any times wanted by user
    if len(args) % 2 != 0:
        print "error: please provide legend and tsung log filename"
        sys.exit(3)

    count   = 0
    legends = []
    files   = []

    for a in args:
        if count % 2 == 0:
            legends.append(a)
        else:
            files.append(a)

        count += 1

    if options.verbose:
        print 'Using %s stats configuration file' % SYS_STATS_CONF

    logs = []
    for logfile in files:
        if not os.access(logfile, os.R_OK):
            print "error: unable to read file %s" % logfile

        else:
            if options.verbose:
                print 'Parsing Tsung log file', logfile
            logs.append((logfile, TsungLog(SYS_STATS_CONF, logfile)))

    if len(logs) != len(args) / 2:
        print 'error while parsing files (%d != %d)' % (len(logs),
                                                        len(args)/2)
        sys.exit(2)

    main(config, logs, legends, options.outdir, options.verbose)
