## ------------------------------------------------------------------------
## iSiLK - A python-based graphical user interface for SiLK
## ------------------------------------------------------------------------
## Copyright (c) 2007-2008 Carnegie Mellon University. All Rights Reserved.
## ------------------------------------------------------------------------
## Authors: Andrew Kompanek (ajk@cert.org)
## ------------------------------------------------------------------------
## GNU General Public License (GPL) Rights pursuant to Version 2, June 1991
## Government Purpose License Rights (GPLR) pursuant to DFARS 252.225-7013
## ------------------------------------------------------------------------
## See the LICENSE file included in this distribution
## ------------------------------------------------------------------------

import wx
import os, os.path, sys, string, tempfile, copy, time, pdb
import constants
import config.silk as config
from analysisresult import change_event_trigger, analysis_job
from analysisresult import AnalysisResult
import datetime
from timeutils import datetime_obj, silk_datetime_range, silk_date, silk_datetime
import timeutils as rave_times
# FIXME - put the format code somewher eelse?
from datacache import formatFieldAsString

#import rave.plugins.times as rave_times
from matplotlib.figure import Figure
import silk

# --------- Graph Results (using MatPlotLib) ---------- #

class GraphResult (AnalysisResult):
    def __init__(self, namePrefix, fileExtension, parent, **kargs):
        AnalysisResult.__init__(self, AnalysisResult.GRAPH_RESULT, 
                    namePrefix, "png", parent, kargs)
        self.fileinfo["command-lines"] = ["(no command line equivalent)"]
        if self.GetParent():
            self.labels = self.GetParent().GetLabels()
        self._xmlChildren =["name", "outputFileName", "fields", "opts", "fileinfo"]
       
    def _xml_postprocess_(self):
        self._changeListeners = []
        self.fields = self.GetParent().GetFields()
           
    def GetRecordCount(self): 
        return self.GetParent().GetRecordCount()
        
    def GetFields(self): 
        return self.GetParent().GetFields()
    
    def IsRemoteFileAvailable(self):
        return False

    def IsLocalFileAvailable(self):
        return True
    
    def GetSize(self):
        return (500, 750)
    
    def OnButtonPress(self, evt):
        pass

    def OnMotion(self, evt):
        pass



from graphutils import RaveFormatter, RaveLogLocator, RaveAutoDateFormatter
from matplotlib.ticker import Formatter, LogLocator, MaxNLocator, \
                              ScalarFormatter, FuncFormatter
from matplotlib.dates import AutoDateLocator, num2date, date2num
import math
from matplotlib.numerix import take, nonzero, absolute
                              
class LogRaveFormatter(ScalarFormatter):
    def __init__(self, y_units=''):
        ScalarFormatter.__init__(self)
        self.y_units = y_units
    
    def get_offset(self):
        return ''
    
    def pprint_val(self, x):
        if x == 0.0:
            return ""
        # 
        label_val = x/10**12
        if label_val >= 1.0:
            return ("%0.1f T%s" % (label_val, self.y_units))
        label_val = x/10**9
        if label_val >= 1.0:
            return ("%0.1f G%s" % (label_val, self.y_units))
        label_val = x/10**6
        if label_val >= 1.0:
            return ("%0.1f M%s" % (label_val, self.y_units))
        label_val = x/10**3
        if label_val >= 1.0: 
            return ("%0.1f k%s" % (label_val, self.y_units))
        if self.orderOfMagnitude >= 0:
            if x >= 1:
                return ("%d %s" % (int(x), self.y_units))
            else:
                return ("%0.2f %s" % (x, self.y_units))
        else:
            return ("%0.2g %s" % (x, self.y_units))
    
    def _set_orderOfMagnitude(self,range):
      #print "Range is %s" % range
      # if scientific notation is to be used, find the appropriate exponent
      # if using an numerical offset, find the exponent after applying the offset
      locs = absolute(self.locs)
      if self.offset: oom = math.floor(math.log10(range))
      else:
          if locs[0] > locs[-1]: val = locs[0]
          else: val = locs[-1]
          if val == 0: oom = 0
          else: oom = math.floor(math.log10(val))
      if oom <= -2:
          self.orderOfMagnitude = oom
      elif oom >= 3:
          self.orderOfMagnitude = oom
      else:
          self.orderOfMagnitude = 0
 

class MplBarGraph:
    def __init__(self, result, figure, y_fields=None, max_rows=10,
                bar_thickness=0.50, yscale='linear', orientation='vertical',
                sort_data=False, title_prefix=""):
                
        valid_x_fields = ["sip", "dip", "sport", "dport", "proto",
            "sensor", "nhip", "date", "scc", "dcc", "sval", "dval"]
        valid_y_fields = ["bytes", "packets", "records",
            "unique_sip", "unique_dip"]
        self.BAR_COLORS = {
            "bytes" : (0.3,0.3,0.9),
            "packets" : (0.9,0.3,0.3),
            "records":(0.3,0.9,0.3),
            "unique_sip":(0.9,0.9,0.3),
            "unique_dip":(0.9,0.3,0.9)}
        self.LABEL_WIDTHS = {
            "sip" : 125.0, "dip" : 125.0, "sport" : 50.0, "dport" : 50.0,
            "proto" : 50.0, "sensor" : 175.0, "nhip" : 125.0
        }
        self.result = result
        self.sort_data = sort_data
        self.figure = figure
        self.title_prefix = title_prefix
        self.max_rows = max_rows
        self.yscale = yscale
        self.orientation = orientation
        if result.GetRecordCount() < max_rows:
            self.max_rows = result.GetRecordCount()
        self.x_fields = [x for x in valid_x_fields if x in result.GetFields()]
        if not y_fields:
            y_fields = result.GetFields()
        self.y_fields = [y for y in valid_y_fields if y in y_fields]
        self.y_fields.reverse()
        # Initialize size constants
        self.graph_count = len(self.y_fields)
        self.BAR_THICKNESS = bar_thickness
        self.MAX_WIDTH=1024
        self.MAX_HEIGHT=1024
        self.hideLabels = False
        self.timeSeriesData = ("date" in self.x_fields)

        if self.orientation == 'vertical':
            if self.max_rows > 25:
                self.px_graph_height = 5.0 * self.max_rows        
            else:
                self.px_graph_height = 45.0 * self.max_rows
            self.px_graph_width = 400.0
        else:
            if self.max_rows > 25:
                self.px_graph_width = 10.0 * self.max_rows 
                self.hideLabels = True # Not enough space
            else:
                self.px_graph_width = 45.0 * self.max_rows
            self.px_graph_height = 200.0
        if self.px_graph_height > self.MAX_HEIGHT:
            self.px_graph_height = self.MAX_HEIGHT
        if self.px_graph_width > self.MAX_WIDTH:
            self.px_graph_width = self.MAX_WIDTH

        self.px_graph_left_border = 85.0
        self.px_graph_right_border = 25.0
        self.px_graph_top_border = 100.0
        if self.orientation == 'vertical':
            self.px_graph_left_border = 150.0
            if "sip" in self.x_fields or "dip" in self.x_fields or "nhip" in self.x_fields:
                self.px_graph_left_border = 225.0
        self.px_graph_yborder = 75.0 
        self.px_height = (self.px_graph_height + self.px_graph_yborder) * self.graph_count + self.px_graph_top_border
        self.px_width = self.px_graph_left_border + self.px_graph_right_border + self.px_graph_width
        self.graph_height = self.px_graph_height / float(self.px_height)
        self.graph_width = self.px_graph_width / float(self.px_width)
        self.yborder = float(self.px_graph_yborder) / float(self.px_height)
        self.top_border = float(self.px_graph_top_border) / float(self.px_height)
        self.left_border = self.px_graph_left_border / float(self.px_width)
        
    # FIXME: more accurate calculation here
    def get_size(self):
        return (self.px_width, self.px_height)
        
    def render(self, render_as_icon = False):
        self._extract_graph_data()
        self._calculate_layout()
        label_data = []
        self.graph_axes = {}
        # Build the labels for each of the xvals
        # (same for all graphs)
        for i in range(self.result.GetRecordCount()): #range(self.max_rows):
            label_values = {}
            for xf in self.x_fields:
                label_values[xf] = self.x_data[xf][i]
            label = self._format_label(label_values)
            if not render_as_icon:
                label_data.append(label)
            else:
                label_data.append("")
        for (position, yf) in enumerate(self.y_fields):
            x_values = []
            # Make a copy of the x_data (as a list of tuples)
            for i in range(0, self.result.GetRecordCount()): #range(self.max_rows):
                xval = []
                for xf in self.x_fields:
                    xval.append(self.x_data[xf][i])
                x_values.append(xval)
            labels = copy.copy(label_data)
            values = copy.copy(self.y_data[yf])
            x_values = copy.copy(x_values)
            if self.sort_data:
                aux = [(x, i) for (i,x) in enumerate(values)]
                aux.sort(reverse=True)
                labels = [labels[i] for (x, i) in aux]
                values = [values[i] for (x, i) in aux]
                x_values = [x_values[i] for (x, i) in aux]
            label_vals = {}
            for xf in self.x_fields:
                label_vals[xf] = xf
            #xlab = self._format_label(label_vals)
            xlab = ", ".join(label_vals)
            if yf == "bytes": title = "Volume in Bytes"
            if yf == "packets": title = "Volume in Packets"            
            if yf == "records": title = "Volume in Flow Counts"
            if yf == "unique_sip": title = "Unique Source IP Addresses"
            if yf == "unique_dip": title = "Unique Destination IP Addresses"
            if self.title_prefix:
                title = "%s: %s\n%s" % (self.title_prefix, title, xlab)
            else:
                title = "%s\n%s" % (title, xlab)                
            # Only graph top max_rows values
            labels = labels[0:self.max_rows]
            values = values[0:self.max_rows]
            # Render it
            ax = self.render_subplot(position, yf, labels, values,
                    "", "", title , render_as_icon=render_as_icon)
            self.graph_axes[ax] = (yf, copy.copy(self.x_fields), x_values, values)
            
    def _format_label(self, label_vals):
        label = ""
        f = self.x_fields
        if "nhip" in f:
            label += str(label_vals["nhip"])
        if "sip" in f:
            label += str(label_vals["sip"])
        if "sport" in f:
            if "sip" in f: label += ":"
            label += str(label_vals["sport"])
        if "dip" in f:
            if "sip" in f or "sport" in f: label += "\n"
            label += str(label_vals["dip"])
        if "dport" in f:
            if "dip" in f: label += ":"
            elif "sip" in f or "sport" in f: label += "\n"
            label += str(label_vals["dport"])
        if "proto" in f:
            label += "\n" + str(label_vals["proto"])
        if "scc" in f:
            label += "\n" + str(label_vals["scc"])
        if "dcc" in f:
            label += "\n" + str(label_vals["dcc"])
        if "sval" in f:
            label += "\n" + str(label_vals["sval"])
        if "dval" in f:
            label += "\n" + str(label_vals["dval"])                                    
                        
        if "sensor" in f:
            if len(label_vals) > 1:
                label += "\n" + str(label_vals["sensor"])                
            else:
                label += " " + str(label_vals["sensor"])
        return label
        
    def _extract_graph_data(self, sort=False):
        ''' Returns x_data and y_data dictionaries '''
        self.x_data = {}; self.y_data = {}
        for xf in self.x_fields:
            self.x_data[xf] = []
        for yf in self.y_fields:
            self.y_data[yf] = []
        for i in range(0, self.result.GetRecordCount()):
            rec = self.result.GetRecordAt(i)
            for xf in self.x_fields:
                if xf == "date":
                    val = date2num(datetime_obj(rec.getfield(xf)))
                else:
                    val = formatFieldAsString(xf, rec.getfield(xf))
                self.x_data[xf].append(val)
            for yf in self.y_fields:
                self.y_data[yf].append(float(rec.getfield(yf)))
                
    def _calculate_layout(self):
        if len(self.y_data.keys()) == 1:
            self.rows = 1; self.cols = 1
        elif len(self.y_data.keys()) == 2:
            self.rows = 2; self.cols = 1
        else: 
            self.rows = len(self.y_data.keys()); self.cols = 1
    
    def get_selection_info(self, ax, x, y):
        ''' Returns a pair of dictionaries, the first mapping from the key
            fields to the values for the category for the bar under the pont
            and a dictionary pointing from the volume field to volume value.'''
        if not ax in self.graph_axes:
            print "Couldn't find axes!!!"
            return
        (yf, x_fields, x_data, y_data) = self.graph_axes[ax]
        closest = None
        for index in range(0, len(x_data)):
            if self.orientation == 'vertical': 
                val_val = x
                key_val = y
            else: 
                val_val = y
                key_val = x
            d = abs(((key_val-self.BAR_THICKNESS*0.5) - (index + 1)))
            if d <= self.BAR_THICKNESS*0.5: 
                closest = index
                break
        if closest == None:
            return None
        if closest >= len(y_data):
            return None
        if val_val > y_data[closest]:
            return None
        return (dict(zip(x_fields, x_data[closest])),
                dict(zip([yf],    [y_data[closest]])))
                
    def get_x_fields(self):
        return self.x_fields
        
    def get_y_fields(self):
        return self.y_fields
        
    def render_subplot(self, pos, field, labels, values, xlab, ylab, title,
        render_as_icon=False):
        graph_rect = [self.left_border, 
                      self.top_border + (self.graph_height+self.yborder)*pos,
                      self.graph_width, self.graph_height]
        count = len(labels)
        key_pos = []
        for i in range(count):
            key_pos.append(i+1)
            
        # Create and format axes
        if self.orientation == 'vertical': 
            xscale = self.yscale; yscale='linear'
        else: 
            yscale=self.yscale; xscale='linear'
        ax = self.figure.add_axes(graph_rect, xscale=xscale, yscale=yscale,
            axis_bgcolor='#dddddd')
        ax.grid(alpha=0.2, color='k', linestyle='-')

        if self.hideLabels and not self.timeSeriesData:
            key_labels = []
            key_ticks = []
            key_locator = None
            key_formatter = None
        elif self.timeSeriesData:
            key_labels = []
            key_ticks = []
            nbins = int(self.px_graph_width / 100.0)
            if nbins == 0: nbins = 1
            key_locator = MaxNLocator(nbins=nbins)
            ax.set_xticklabels(key_labels, horizontalalignment='center',
                fontname='Sans', rotation=0, size=12)
            class SimpleFormatter(ScalarFormatter):
                def __init__(self, graph):
                    ScalarFormatter.__init__(self)
                    self.graph = graph
                def __call__(self, x, pos=0):
                    if x >= self.graph.max_rows: return ''
                    date = num2date(self.graph.x_data["date"][int(x)])
                    date_str = "%02d:%02d\n%d/%d/%d" % (date.hour, date.minute,
                        date.year, date.month, date.day)
                    return date_str
                    #return silk_datetime(num2date(val))
                    return 'test'
            key_formatter = SimpleFormatter(self)            
        else:
            key_labels = [""]
            key_locator = None
            key_formatter = None
            key_labels.extend(labels)
            key_labels.append("")
            key_ticks = [x+self.BAR_THICKNESS*0.5 for x in range(0,count+2)]            
            ax.set_xlim(xmin=0.5, xmax=len(key_pos)+1.25)
        
        if self.orientation == 'vertical':
            (key_axis, val_axis) = (ax.yaxis, ax.xaxis)
            ax.set_yticks(key_ticks)  
            ax.set_ylim(ymin=0.5, ymax=len(key_pos)+1.25)
            ax.set_yticklabels(key_labels, horizontalalignment='right',
                fontname='Sans', rotation=0, size=12)    
            ax.set_xlabel(ylab)
            ax.set_ylabel(xlab)      
        else:
            (key_axis, val_axis) = (ax.xaxis, ax.yaxis) 
            if not self.timeSeriesData:           
                ax.set_xticks(key_ticks)    
                ax.set_xticklabels(key_labels, horizontalalignment='center',
                    fontname='Sans', rotation=0, size=12)
            ax.set_xlabel(ylab)
            ax.set_ylabel(xlab)
        # Set formatting for axes
        if self.yscale == 'log':
            val_locator = RaveLogLocator()
            val_formatter = LogRaveFormatter(y_units='')            
        else:
            val_formatter = RaveFormatter(y_units='') 
            val_locator = MaxNLocator(nbins=4)
        val_axis.set_major_formatter(val_formatter)
        val_axis.set_major_locator(val_locator)

        if key_locator:
            key_axis.set_major_locator(key_locator)
        if key_formatter:
            key_axis.set_major_formatter(key_formatter)                
        
        # FIXME -- How can I make this work for my bar graphs?
        
        #key_locator = AutoDateLocator()
        #key_axis.set_major_locator(key_locator)
        #key_axis.set_major_formatter(RaveAutoDateFormatter(key_locator))
        
        # Render the bar graph
        if self.orientation == 'vertical':
            bar = ax.barh(bottom=key_pos, width=values, height=self.BAR_THICKNESS,
                color=self.BAR_COLORS[field])
        else:
            bar = ax.bar(left=key_pos, height=values, width=self.BAR_THICKNESS,
                color=self.BAR_COLORS[field]) #, alpha=0.75)
        if render_as_icon:
            ax.set_xticklabels([])
            ax.set_xticks([])
        if not render_as_icon:
            ax.set_title(title+"\n")
            ax.set_xlabel(ylab)
            ax.set_ylabel(xlab)
        return ax


class CountGraph (GraphResult, MplBarGraph):
    def __init__(self, parent=None, name=None, y_fields=None, max_rows=10,
                    bar_thickness=0.75, orientation='vertical', yscale='log',
                    sort_data=True, title_prefix=None):
        GraphResult.__init__(self, name, "png", parent, y_fields=y_fields,
            max_rows = max_rows, bar_thickness=bar_thickness, 
            orientation=orientation, yscale=yscale, sort_data = sort_data,
            title_prefix=title_prefix)
            
        print "Creating CountGraph with orientation=%s" % orientation
        self.name = name
        if self.name == None:
            self.name = "Uniq Bar Graph"
        self._selection = None
        self.y_fields = y_fields
        self.max_rows = max_rows
        self.title_prefix = title_prefix
        self.yscale = yscale
        self.bar_thickness = bar_thickness
        self.orientation = orientation
        self.sort_data = sort_data
           
    def _xml_postprocess_(self):
        # FIXME: integrate graphs better into framework for serialization
        self.y_fields = self.opts['y_fields']
        self.max_rows = self.opts['max_rows']
        self.yscale   = self.opts['yscale']
        self.bar_thickness = self.opts['bar_thickness']
        self.orientation = self.opts['orientation']
        self.sort_data = self.opts['sort_data']
        GraphResult._xml_postprocess_(self)
        
    def RunAnalysis(self):
        self._UpdateStatus(constants.AVAILABLE_LOCAL, "Loaded")
        
    def GetFields(self):
        fields = copy.copy(self.graph.get_x_fields())
        fields.extend(self.graph.get_y_fields())
        return fields
        
    def _getDrillDownOptions(self, key_map):
        opts = copy.copy(key_map)
        if "sip" in opts:
            opts["saddress"] = opts["sip"]
            del opts["sip"]
        if "dip" in opts:
            opts["daddress"] = opts["dip"]
            del opts["dip"]
        if "sval" in opts:
            opts["pmap-saddress"] = "sval"
            del opts["sval"]
        if "dval" in opts:
            opts["pmap-daddress"] = "dval"            
            del opts["dval"]            
        if "date" in opts:
            start = num2date(opts["date"])
            bin_size = self.GetParent().opts["bin-size"]
            end = start + datetime.timedelta(seconds=bin_size)
            val = silk_datetime_range(start,end)
            del opts["date"]
            opts["stime"] = val
        dd_opts = {}
        for o in ["daddress", "saddress", "scc", "dcc", "sval", "dval",
            "sport", "dport", "proto", "stime"]:
            if o in opts:
                dd_opts[o] = opts[o]
        return dd_opts
        
        
    def GetSelection(self):
        if not self._selection: return None
        (key_map, val_map) = self._selection
        sel = copy.copy(key_map)
        sel.update(val_map)
        return sel
            
    def GetSelectionString(self):
        print "In GetSelectionString"
        if not self._selection: return ""
        (key_map, val_map) = self._selection
        
        opts = self._getDrillDownOptions(key_map)
        s = ""  
        for o in opts: 
            s += "%s=%s " % (o, opts[o])
        print "Selection is %s" % s
        return s
            
    def OnDoubleClick(self, evt):
        if self._selection:
            (key_map, val_map) = self._selection
            opts = self._getDrillDownOptions(key_map)
            parent = wx.GetApp().GetSelectedFilterData()
            wx.GetApp().refine(parent=parent, **opts)
         
    def OnButtonPress(self, evt):
        # Find the bar that's closest to where they clicked
        if evt.inaxes:
            print "In axes"
            sel = self.graph.get_selection_info(evt.inaxes, evt.xdata, evt.ydata)
            if sel == None: return
            (key_map, val_map) = sel
            if evt.button == 3:
                # make a menu
                menu = wx.Menu()
                # add some other items
                options = ""
                # FIXME: convert field names to rwfilter options 
                # create a utility
                for opt in key_map: 
                    options += " --%s=%s" % (opt, key_map[opt])
                
                menu.Append(9001, "Run rwfilter%s" % options)
                # make a submenu
                #sm = wx.Menu()
                #sm.Append(-1, "sub item 1")
                #menu.AppendMenu(-1, "Test Submenu", sm)
                wx.GetApp().GetTopWindow().PopupMenu(menu)

    
    def OnMotion(self, evt):
        # Don't propagate exceptions within this handler to avoid very bad
        # feedback to the user.  Instead siimply don't update the selection.
        try:
            if evt.inaxes:
                self._selection = self.graph.get_selection_info(evt.inaxes, evt.xdata, evt.ydata)
                if not self._selection: return
                (key_map, val_map) = self._selection
                opts = self._getDrillDownOptions(key_map)
        except:
            return

    def Plot(self, fig):
        def make_canvas():
            try:
                from matplotlib.backends.backend_agg import FigureCanvasAgg
                fig_canvas = FigureCanvasAgg
            except:
                from matplotlib.backends.backend_cairo import FigureCanvasCairo
                fig_canvas = FigureCanvasCairo
            self.fig_canvas = fig_canvas
            return fig_canvas
            
        if fig == None:
            dpi = 75.0
            (width, height) = self.GetSize()
            in_w = width / dpi
            in_h = height / dpi
            fig = Figure(figsize=(in_w, in_h), dpi=dpi) #**kargs
            graph = self._create_graph(fig)
            graph.render()
            out_file_name = self.GetLocalPath()
            canvas = make_canvas()(fig)
            fig.savefig(out_file_name, dpi=dpi)
        else:
            self.graph = self._create_graph(fig)
            self.graph.render()
        
    def _create_graph(self, fig):
        graph = MplBarGraph(self.GetParent(), figure=fig,
            y_fields=self.y_fields,
            max_rows=self.max_rows, orientation=self.orientation, 
            bar_thickness=self.bar_thickness, yscale=self.yscale,
            sort_data=self.sort_data, title_prefix=self.title_prefix)
        return graph
        
    def GetSize(self):
        graph = self._create_graph(None)
        return graph.get_size()
                    