## ------------------------------------------------------------------------
## 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
## ------------------------------------------------------------------------

"""
Matplotlib visualizations
"""
import sys

# vvv THIS IS FIRST SO IT CAN BE SURE TO PATCH THE font_manager FIRST

# HACK
# Make font_manager.FontProperties do something sane.

font_filename = None
font_small_size = 10
font_large_size = 20

import matplotlib.font_manager
OldFontProperties = matplotlib.font_manager.FontProperties
class NewFontProperties(OldFontProperties):
    def __init__(self, *args, **kwargs):
        if ('fname' not in kwargs) and (font_filename is not None):
            kwargs['fname'] = font_filename
        OldFontProperties.__init__(self, *args, **kwargs)
matplotlib.font_manager.FontProperties = NewFontProperties

def set_font_filename(filename):
    global font_filename
    font_filename = filename

# ^^^ THIS IS FIRST SO IT CAN BE SURE TO PATCH THE font_manager FIRST

import sys

#from dataset import Dataset, seq_set, seq_find
from timeutils import datetime_obj, bin_datetime
from timeutils import moving_average

from datetime import datetime, timedelta
from calendar import timegm

from threading import Lock

from matplotlib.ticker import Formatter, LogLocator, MaxNLocator, \
                              ScalarFormatter, FuncFormatter
from matplotlib.figure import Figure
from matplotlib.axes import Axes
from matplotlib.dates import AutoDateLocator, date2num, num2date
from matplotlib.numerix import take, nonzero, absolute
from matplotlib.font_manager import FontProperties
from matplotlib.artist import setp
from matplotlib.text import Text
from matplotlib.lines import Line2D
from matplotlib.patches import Polygon, Rectangle

fig_canvas = None
try:
    from rave.plugins.mpl_kludge_cairo import FigureCanvasCairo
    fig_canvas = FigureCanvasCairo
except:
    try:
        from matplotlib.backends.backend_agg import FigureCanvasAgg
        fig_canvas = FigureCanvasAgg
    except:
        from matplotlib.backends.backend_cairo import FigureCanvasCairo
        fig_canvas = FigureCanvasCairo

import re, locale, time, math

renderer_lock = Lock()

strftime_dt_illegal_s = re.compile(r"((^|[^%])(%%)*%s)")

# Works for wider range of dates than normal
def strftime_dt(dt, fmt):
    
    def findall(self, text, substr):
        # Also finds overlaps
        sites = []
        i = 0
        while 1:
            j = text.find(substr, i)
            if j == -1:
                break
            sites.append(j)
            i=j+1
        return sites
    
    fmt = strftime_dt_illegal_s.sub(r"\1", fmt)
    fmt = fmt.replace("%s", "s")
    if dt.year > 1900:
            return unicode(time.strftime(fmt, dt.timetuple()), locale.getpreferredencoding())

    year = dt.year
    # For every non-leap year century, advance by
    # 6 years to get into the 28-year repeat cycle
    delta = 2000 - year
    off = 6*(delta // 100 + delta // 400)
    year = year + off

    # Move to around the year 2000
    year = year + ((2000 - year)//28)*28
    timetuple = dt.timetuple()
    s1 = time.strftime(fmt, (year,) + timetuple[1:])
    sites1 = findall(s1, str(year))

    s2 = time.strftime(fmt, (year+28,) + timetuple[1:])
    sites2 = findall(s2, str(year+28))

    sites = []
    for site in sites1:
        if site in sites2:
            sites.append(site)

    s = s1
    syear = "%4d" % (dt.year,)
    for site in sites:
        s = s[:site] + syear + s[site+4:]

    return unicode(s, locale.getpreferredencoding())

class RaveAutoDateFormatter(Formatter):
    def __init__(self, locator):
        self.locator = locator

    def __call__(self, x, pos=0):
        scale = float(self.locator._get_unit())
        (dmin, dmax) = self.locator.datalim_to_dt()
        dt = num2date(x)
        dmin_n = date2num(dmin)
        dmax_n = date2num(dmax)

        # Format for top line: shown for every tick
        fmt1 = "%H:%M"                  # 1:00

        # Format for bottom line: shown only when the value changes
        fmt2 = "%b %e, %Y"              # Jan 1, 1970

        if scale == 365.0:              # Tick on years
            fmt1 = "%Y"
            fmt2 = ""
        elif scale == 30.0:             # Tick on months
            fmt1 = "%B"
            fmt2 = "%Y"
        elif scale == 1.0 or scale == 7.0: # Weeks or days
            fmt1 = "%e"
            fmt2 = "%b %Y"
        elif scale == 1.0/24.0:         # Hours
            fmt1 = "%H:%M"
            fmt2 = "%b %e, %Y"

        # Lesser time scales handled by the defaults

        dmin_s = strftime_dt(dmin, fmt2)
        dmax_s = strftime_dt(dmax, fmt2)
        if dmin_s == dmax_s:
            # No change in display range:
            if x < dmin_n + scale or x > dmax_n - scale:
                # Display at start and end of range
                return strftime_dt(dt, fmt1 + "\n" + fmt2)
            else:
                return strftime_dt(dt, fmt1)
        else:
            # Change occurs within display range:
            dt_pre_s = strftime_dt(num2date(x - scale), fmt2)
            dt_s = strftime_dt(dt, fmt2)
            if dt_pre_s <> dt_s:
                # Display at points of change
                return strftime_dt(dt, fmt1 + "\n" + fmt2)
            else:
                return strftime_dt(dt, fmt1)

# Log Locator that does not display a tick at 10^0 (1)
class RaveLogLocator(LogLocator):
    def __call__(self):
        ticklocs = LogLocator.__call__(self)
        return take(ticklocs, nonzero(ticklocs<>1.0))

class RaveFormatter(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 ""
        if self.orderOfMagnitude >= 12:
            return ("%0.1f T%s" % (x/10**12, self.y_units))
        elif self.orderOfMagnitude >= 9:
            label_val = x/10**9
            if label_val < 1.0:
                label_val = x/10**6
                return ("%0.1f M%s" % (label_val, self.y_units))
            return ("%0.1f G%s" % (label_val, self.y_units))
        elif self.orderOfMagnitude >= 6:
            return ("%0.1f M%s" % (x/10**6, self.y_units))
        elif self.orderOfMagnitude >= 3:
            return ("%0.1f k%s" % (x/10**3, self.y_units))
        elif 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

# For working with "zero" on log scales.
def zeros_to_ones(x):
    if x == 0.0:
        return 1.0
    return x

# For abbreviating raw counts
# TODO: this probably belongs somewhere else....
def format_count(count):
    kb = 1024.0
    mb = 1024.0 * kb
    gb = 1024.0 * mb
    tb = 1024.0 * gb

    if count >= tb:
        return "%.2fT" % (count/tb)
    if count >= gb:
        return "%.2fG" % (count/gb)
    if count >= mb:
        return "%.2fM" % (count/mb)
    if count >= kb:
        return "%.2fk" % (count/kb)
    return "%.2fb" % count

