#!/usr/bin/python
#  Copyright (C) 2012 by Carnegie Mellon University.
#  
#  @OPENSOURCE_HEADER_START@
#  Use of the Rayon and related source code is subject to the terms of
#  the following licenses:
#  
#  GNU Public License (GPL) Rights pursuant to Version 2, June 1991
#  Government Purpose License Rights (GPLR) pursuant to DFARS 252.227.7013
#  
#  NO WARRANTY
#  
#  ANY INFORMATION, MATERIALS, SERVICES, INTELLECTUAL PROPERTY OR OTHER 
#  PROPERTY OR RIGHTS GRANTED OR PROVIDED BY CARNEGIE MELLON UNIVERSITY 
#  PURSUANT TO THIS LICENSE (HEREINAFTER THE "DELIVERABLES") ARE ON AN 
#  "AS-IS" BASIS. CARNEGIE MELLON UNIVERSITY MAKES NO WARRANTIES OF ANY 
#  KIND, EITHER EXPRESS OR IMPLIED AS TO ANY MATTER INCLUDING, BUT NOT 
#  LIMITED TO, WARRANTY OF FITNESS FOR A PARTICULAR PURPOSE, 
#  MERCHANTABILITY, INFORMATIONAL CONTENT, NONINFRINGEMENT, OR ERROR-FREE 
#  OPERATION. CARNEGIE MELLON UNIVERSITY SHALL NOT BE LIABLE FOR INDIRECT, 
#  SPECIAL OR CONSEQUENTIAL DAMAGES, SUCH AS LOSS OF PROFITS OR INABILITY 
#  TO USE SAID INTELLECTUAL PROPERTY, UNDER THIS LICENSE, REGARDLESS OF 
#  WHETHER SUCH PARTY WAS AWARE OF THE POSSIBILITY OF SUCH DAMAGES. 
#  LICENSEE AGREES THAT IT WILL NOT MAKE ANY WARRANTY ON BEHALF OF 
#  CARNEGIE MELLON UNIVERSITY, EXPRESS OR IMPLIED, TO ANY PERSON 
#  CONCERNING THE APPLICATION OF OR THE RESULTS TO BE OBTAINED WITH THE 
#  DELIVERABLES UNDER THIS LICENSE.
#  
#  Licensee hereby agrees to defend, indemnify, and hold harmless Carnegie 
#  Mellon University, its trustees, officers, employees, and agents from 
#  all claims or demands made against them (and any related losses, 
#  expenses, or attorney's fees) arising out of, or relating to Licensee's 
#  and/or its sub licensees' negligent use or willful misuse of or 
#  negligent conduct or willful misconduct regarding the Software, 
#  facilities, or other rights or assistance granted by Carnegie Mellon 
#  University under this License, including, but not limited to, any 
#  claims of product liability, personal injury, death, damage to 
#  property, or violation of any laws or regulations.
#  
#  Carnegie Mellon University Software Engineering Institute authored 
#  documents are sponsored by the U.S. Department of Defense under 
#  Contract FA8721-05-C-0003. Carnegie Mellon University retains 
#  copyrights in all material produced under this contract. The U.S. 
#  Government retains a non-exclusive, royalty-free license to publish or 
#  reproduce these documents, or allow others to do so, for U.S. 
#  Government purposes only pursuant to the copyright license under the 
#  contract clause at 252.227.7013.
#  @OPENSOURCE_HEADER_END@

from rayon import miscutils, data
from rayon.appsupport import trendscript
from rayon.cliutils import mkcb_color

from netsa.data import format

from optparse import OptionParser, OptionValueError, SUPPRESS_HELP
import pprint
import os
import sys

# all() is present in 2.5+, but not in 2.4, so:
try:
    all
except NameError:
    def all(i):
        for _ in i:
            if not _: return False
        return True


#
# Argument processing
#


defaults = {
    'input_path'       : '-',
    'output_path'      : None,
    'width'            : 600,
    'height'           : 800,
    'title'            : "",
    'background_color' : miscutils.parse_color_string("#ffffffff"),
    'top_color'        : miscutils.parse_color_string("#ff000080"),
    'btm_color'        : miscutils.parse_color_string("#0000ff80"),
    'title_width'      : 100,
    "bytes"            : False,
    "packets"          : False,
    "records"          : False,
    "no_units_set"     : True,
    'print_options'    : False,
}

def parse_options():
    op = OptionParser(usage="%prog [options] - "
                      "generate strip plots of multiple time series")
    op.set_defaults(**defaults)
    op.add_option("--input-path", dest="input_path",
                  help="input file ('-' == STDIN)")
    op.add_option("--output-path", dest="output_path",
                  help="(REQUIRED) output path")
    op.add_option("--width", type="int", dest="width",
                  help="image width (pixels for PNG, otherwise points) "
                  "(Default: %d)" % defaults['width'])
    op.add_option("--height", type="int", dest="height",
                  help="image height (pixels for PNG, otherwise points) "
                  "(Default: %d)" % defaults['width'])
    op.add_option("--title", type="str", dest="title",
                  help="title of chart (ignored for single series)")
    op.add_option('--background-color', metavar="COLOR",
                  action="callback", callback=mkcb_color("background_color"),
                  dest="background_color", type="str",
                  help="Color of background (Default: \"#FFFFFFFF\" (white))")
    op.add_option('--top-color', metavar="COLOR",
                  action="callback", callback=mkcb_color("top_color"),
                  dest="top_color", type="str",
                  help="Color of upper time series "
                  "(Default: \"#FF000080\" (light red))")
    op.add_option('--bottom-color', metavar="COLOR",
                  action="callback", callback=mkcb_color("btm_color"),
                  dest="top_color", type="str",
                  help="Color of lower time series "
                  "(Default: \"#0000FF80\" (light blue))")
    op.add_option('--bytes', action="store_true", dest="bytes",
                  help="Indicate input is bytes")
    op.add_option('--packets', action="store_true", dest="packets",
                  help="Indicate input is packets")
    op.add_option('--records', action="store_true", dest="records",
                  help="Indicate input is records (flows)")
    op.add_option('--print-options', action="store_true", dest="print_options",
                  help=SUPPRESS_HELP)
    options, args = op.parse_args()
    return options

#
# Debugging
#

def pprint_opts(opts):
    pp = pprint.PrettyPrinter()
    pp.pprint(opts.__dict__)


#
# Building the visualization
#



class ConfigurationError(Exception):
    pass

class InputError(Exception):
    pass

def sanity_check_io_opts(opts):
    if opts.input_path is None:
        if sys.stdin.isatty():
            raise ConfigurationError("--input-path was not supplied, and "
                                     "STDIN is not a pipe")
        else:
            opts.input_path = "-"

    if opts.output_path is None:
        raise ConfigurationError("--output-path was not supplied")
        

def sanity_check_unit_opts(opts):
    num_set = 0
    if opts.bytes   : num_set += 1
    if opts.packets : num_set += 1
    if opts.records : num_set += 1

    if num_set > 1:
        raise ConfigurationError("Only one of --bytes, --packets or --records "
                                 "may be set")

    if num_set == 1:
        opts.no_units_set = False


def main():
    opts = parse_options()
    if opts.print_options:
        pprint_opts(opts)
        return
    
    sanity_check_io_opts(opts)
    sanity_check_unit_opts(opts)

    if opts.input_path == "-":
        in_data = data.Dataset.from_stream(sys.stdin)
    else:
        in_data = data.Dataset.from_file(opts.input_path)


    # Data should look like this:
    #
    # #date|web_in|web_out|mail_in|mail_out
    # 2010/01/01:00|100|200|44|898

    column_names = in_data.get_column_names()
    if len(column_names) == 0:
        column_names.append('_')
        # Come up with (unsatisfying) defaults
        column_names.extend(["Series %d_%s" % (i+1, d)
                             for i in xrange(in_data.get_num_cols())
                             for d in ("in", "out")])
                            
    # We don't care what the first column is named
    series_names_raw = column_names[1:]
    # Populate input for stripplot
    titles  = []
    tops    = []
    bottoms = []
    while len(series_names_raw) != 0:
        try:
            name_1, name_2 = (series_names_raw.pop(0),
                            series_names_raw.pop(0))
        except IndexError:
            raise InputError("Improperly formatted column names: %s"
                             % str(column_names))
    
        series_1, dir_1 = name_1.split("_")
        dir_1 = dir_1.lower()
        if dir_1 not in ("in", "out"):
            raise InputError("Improperly formatted column names for series %s"
                             % series_1)

        series_2, dir_2 = name_2.split("_")
        if series_2 != series_1:
            raise InputError("Improperly formatted column names for series %s:"
                             " both _in and _out columns are required"
                             % series_1)
        dir_2 = dir_2.lower()
        if dir_2 not in ("in", "out"):
            raise InputError("Improperly formatted column names for series %s"
                             % series_1)
        if dir_1 == dir_2:
            raise InputError("Improperly formatted column names for series %s:"
                             " need exactly one of '_in' and '_out'"
                             % series_1)
        if dir_1 == "out":
            # Swap ordering
            old_name_1 = name_1
            name_1     = name_2
            name_2     = old_name_1
            
        titles.append(series_1)
        tops.append(in_data.get_column(name_1))
        bottoms.append(in_data.get_column(name_2))

    if len(titles) == 0:
            raise InputError("No series to plot")

    # Annotation function
    def annotate(row):
        datum = row[1]
        if opts.bytes:
            return format.num_prefix(datum, "b")
        elif opts.packets:
            return format.num_prefix(datum, "p")
        elif opts.records:
            return format.num_prefix(datum, "f")
        return format.num_exponent(datum)

    try:
        extension = opts.output_path[opts.output_path.rindex(".") + 1:]
    except ValueError:
        raise ConfigurationError("Output path must have extension")

    fobj = open(opts.output_path, 'w')

    if len(titles) == 1:
        trendscript.stripplot_single(binned_time=in_data.get_column(0),
                                     top=tops[0],
                                     bottom=bottoms[0],
                                     title=titles[0],
                                     fobj=fobj,
                                     output_type=extension,
                                     width=opts.width,
                                     height=opts.height,
                                     padding=(10, 10, 10, 10),
                                     title_width=opts.title_width,
                                     bgcolor=opts.background_color,
                                     top_color=opts.top_color,
                                     bottom_color=opts.btm_color,
                                     annotation_function=annotate)
    else:
        trendscript.stripplot_multi(binned_time=in_data.get_column(0),
                                    tops=tops,
                                    bottoms=bottoms,
                                    titles=titles,
                                    fobj=fobj,
                                    output_type=extension,
                                    width=opts.width,
                                    height=opts.height,
                                    vertical_padding=40,
                                    horizontal_padding=10,
                                    title_width=opts.title_width,
                                    bgcolor=opts.background_color,
                                    top_color=opts.top_color,
                                    bottom_color=opts.btm_color,
                                    annotation_function=annotate,
                                    chart_title=opts.title)
                         
        
try:
    main()
except KeyboardInterrupt:
    print "Application interrupted by user"
    sys.exit(1)
except ConfigurationError, e:
    print "Configuration error: %s" % str(e)
    print "Type 'rystripplot --help' for more information"
    sys.exit(1)
except InputError, e:
    print "Input error: %s" % str(e)
    print "Type 'rystripplot --help' for more information"
    sys.exit(1)
except Exception, e:
    if os.getenv("RY_PDB") is not None:
        print str(e)
        import pdb
        pdb.post_mortem(sys.exc_info()[2])
        raise
    else:
        raise
