######################################################################
# Copyright (c) 2007, Petteri Aimonen <jpa@kapsi.fi>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice and this list of conditions.
#     * Redistributions in binary form must reproduce the above copyright
#       notice and this list of conditions in the documentation and/or other
#       materials provided with the distribution.

'''Graphpanels and animal data plotting'''

import wx
import events
import datetime
import settings
from dataaccess import animal_getfeeds, animals_getmaxfeeds, animals_in_section
from dataaccess import groups_in_section, animals_in_group, sectionlist
from dataaccess import group_aggfeeds, section_aggfeeds, groupfeed_average
from dataaccess import getgroup, groupfeed_animalcount, animalcount
from dataaccess import pens_per_valve
from matplotlib.dates import date2num
import matplotlib.dates as mpldates
from matplotlib.font_manager import FontProperties
from mplutils import PlotPanel
from pymisc import _

def aggregate_data(fromdate, todate, section = None, group = None):
    '''Get aggregate functions (xdata, min, max, avg) for group or
    for section. Either has to be given.'''
    
    if section is not None:
        aggdata = section_aggfeeds(section, fromdate, todate)
        groups = groups_in_section(section)
    elif group is not None:
        aggdata = group_aggfeeds(group, fromdate, todate)
        groups = [group]
    else:
        raise ValueError, "Either group or section must be given."
    
    xdata = range(date2num(fromdate), date2num(todate) + 1)
    mindata = [None] * len(xdata)
    maxdata = [None] * len(xdata)
    avgdata = [None] * len(xdata)
    
    for date, mingrams, maxgrams, avggrams in aggdata:
        index = int(date2num(date) - xdata[0])
        mindata[index] = mingrams / 1000.
        maxdata[index] = maxgrams / 1000.
        avgdata[index] = avggrams / 1000.
    
    return xdata, mindata, maxdata, avgdata

def setxticks(subplot, fromdate, todate):
    '''Setup x tick labels and grid for subplot.'''
    wlocator = mpldates.WeekdayLocator(byweekday = (mpldates.MO, mpldates.FR))
    subplot.xaxis.set_major_locator(wlocator)
    subplot.xaxis.set_minor_locator(mpldates.DayLocator())
    
    subplot.set_xlim([date2num(fromdate), date2num(todate)])
    subplot.set_xlabel(_("Day of month"))
    
    subplot.grid(True)
    
    if (todate - fromdate).days > 20: # how many labels will fit?
        subplot.xaxis.set_major_formatter(mpldates.DateFormatter("%d"))
        subplot.xaxis.set_minor_formatter(mpldates.DateFormatter(""))
    else:
        subplot.xaxis.set_major_formatter(mpldates.DateFormatter(""))
        subplot.xaxis.set_minor_formatter(mpldates.DateFormatter("%d"))

def plot_nodata(subplot, fromdate, todate):
    '''Display graph with no data'''
    subplot.plot_date([date2num(fromdate), date2num(todate)], [None, None])
    setxticks(subplot, fromdate, todate)

def has_data(data):
    '''Returns True if there are non-None values in data.'''
    for d in data:
        if d is not None:
            return True
    return False

def plot_animal(subplot, animalid, fromdate, todate, yMax = None):
    '''Plot one animal to a subplot, filling in subplot.userdata'''
    subplot.userdata = ('animal', animalid)
    subplot.set_title(_("Animal %d") % animalid)
    subplot.set_ylabel(_("kg"))
    
    data = animal_getfeeds(animalid, fromdate, todate)
    
    if not data:
        plot_nodata(subplot, fromdate, todate)
        return
    
    xdata = range(date2num(fromdate), date2num(todate) + 1)
    ydata = [None] * len(xdata) # Leave None if no data available -> gap in line
    
    for date, grams in data:
        index = int(date2num(date) - xdata[0])
        ydata[index] = grams / 1000.
    
    if has_data(ydata):
        subplot.plot_date(xdata, ydata, 'bo-', xdate = True, markersize = 3)

    if yMax is None:
        yMax = max(ydata) + 0.5

    subplot.set_ylim([-0.1, yMax])
    setxticks(subplot, fromdate, todate)

def plot_many(subplot, fromdate, todate, yMax = None,
              section = None, group = None):
    '''Plot multi-animal plot to subplot, with min/max/avg curves.
    Either section or group has to be given.'''
    subplot.set_ylabel(_("kg"))
    
    aggdata = aggregate_data(fromdate, todate, section = section, group = group)
    xdata, mindata, maxdata, avgdata = aggdata
    
    if not xdata:
        plot_nodata(subplot, fromdate, todate)
        return
    
    if has_data(mindata) and has_data(maxdata) and has_data(avgdata):
        subplot.plot_date(xdata, mindata, 'bo-', xdate = True, markersize = 3)
        subplot.plot_date(xdata, maxdata, 'ro-', xdate = True, markersize = 3)
        subplot.plot_date(xdata, avgdata, 'go-', xdate = True, markersize = 3)
    
    if yMax is None:
        yMax = max(maxdata) + 0.5

    subplot.set_ylim([-0.1, yMax])
    setxticks(subplot, fromdate, todate)

def plot_groupfeed(subplot, groupid, fromdate, todate, yMax = None):
    '''Plot groupfeed curves for group, fill in subplot.userdata'''
    subplot.userdata = None
    subplot.set_title(_("Groupfeed group %d") % groupid)
    subplot.set_ylabel(_("kg"))
    
    data = groupfeed_average([groupid], fromdate, todate)
    xdata = [date2num(d) for d, g in data]
    ydata = [g for d, g in data]
    
    for i in range(len(ydata)):
        if ydata[i] is not None:
            ydata[i] /= 1000.
    
    subplot.plot_date(xdata, ydata, 'mo-', xdate = True, markersize = 3)
    
    if yMax is None:
        yMax = max(ydata) + 0.5

    subplot.set_ylim([-0.1, yMax])
    setxticks(subplot, fromdate, todate)

def plot_group(subplot, groupid, fromdate, todate, yMax = None):
    '''Plot one group to a subplot, filling in subplot.userdata'''
    subplot.userdata = ('group', groupid)
    subplot.set_title(_("Group %d") % groupid)
    
    plot_many(subplot, fromdate, todate, yMax, group = groupid)

def plot_section(subplot, section, fromdate, todate, yMax = None):
    '''Plot one section to a subplot, filling in subplot.userdata'''
    subplot.userdata = ('section', section)
    subplot.set_title(_("Section %d") % section)
    
    plot_many(subplot, fromdate, todate, yMax, section = section)

class MultiGraph(wx.ScrolledWindow):
    '''Multiple graph plotting'''
    
    cols = 2
    
    def __init__(self, *args, **kwargs):
        wx.ScrolledWindow.__init__(self, *args, **kwargs)
        
        # We can't subclass plotpanel because we want to act like ScrolledWindow
        self.plotpanel = PlotPanel(self)
        self.canvas = self.plotpanel.canvas
        self.figure = self.plotpanel.figure
        self.plotpanel.draw = self.draw
        
        # self.graphs is list of tuples (function, list(args), dict(kwargs))
        # Items are called like Function(subplot, *args, **kwargs)
        self.active = False
        self.graphs = []
        self.infotext = ""
        self.showlegend = False
        
        self.cols = settings.parser.getint("GUI", "graphcols")
        
        self.canvas.mpl_connect('motion_notify_event', self.OnMPLMove)
        self.canvas.mpl_connect('button_press_event', self.OnMPLClick)
    
        self.Bind(wx.EVT_SIZE, self.OnSize)
        self.canvas.Bind(wx.EVT_MOUSEWHEEL, self.OnScroll)
        
        self.oldsize = None
    
    wheelScroll = 0
    def OnScroll(self, evt):
        '''Handle mouse scrolling.
        TODO: find way to tell this directly to ScrolledWindow'''
        delta = evt.GetWheelDelta()
        rot = evt.GetWheelRotation()
        linesPer = evt.GetLinesPerAction()
        ws = self.wheelScroll
        ws = ws + rot
        lines = ws / delta
        ws = ws - lines * delta
        self.wheelScroll = ws
        if lines != 0:
            lines = lines * linesPer
            vsx, vsy = self.GetViewStart()
            scrollTo = max(0, vsy - lines)
            self.Scroll(-1, scrollTo)
    
    def OnMPLMove(self, evt):
        self.SetFocus() # Scrolling needs focus on MSW
        
        if evt.inaxes and evt.inaxes.userdata is not None:
            cursor = wx.StockCursor(wx.CURSOR_HAND)
        else:
            cursor = wx.StockCursor(wx.CURSOR_DEFAULT)
        
        self.canvas.SetCursor(cursor)
    
    def OnMPLClick(self, mplevent):
        if not mplevent.inaxes:
            return
        
        wxevent = events.NodeSelectEvent()
        wxevent.userdata = mplevent.inaxes.userdata
        
        if wxevent.userdata is not None:
            wx.PostEvent(self, wxevent)
    
    def OnSize(self, evt = None):
        if self.GetSize() != self.oldsize:
            self.oldsize = self.GetSize()
            self.Redraw()
    
    def AdjustSize(self):
        '''Adjust scroll size & col/row counts. Set self.active'''
        w, h = self.GetSize()
        print w
        if w < 100: # Some initial too small size
            self.active = False
            return
        
        if not self.graphs:
            self.active = False
            return
        
        # Calculate number of rows
        self.rows = len(self.graphs) // self.cols
        if len(self.graphs) % self.cols:
            self.rows += 1
        
        # Determine if we should scroll
        if h // self.rows < 150:
            h = self.rows * 200
            self.SetScrollRate(0, 20)
        else:
            self.SetScrollRate(0, 0)
        
        self.SetVirtualSize((-1, h))
        w, h = self.GetVirtualSize()
        self.plotpanel.SetSize((w, h))
        
        # Matplotlib uses scale 0..1 in args -> converting from pixels
        sw = w / self.cols # Subgraph size
        sh = h / self.rows
        self.figure.subplots_adjust(left = 60. / w, right = 1. - 10. / w,
                                    top = 1. - 30. / h, bottom = 50. / h,
                                    wspace = 50. / sw, hspace = 40. / sh)
        
        self.active = True
    
    def Redraw(self):
        '''Public function to force a redraw'''
        self.AdjustSize()
        # Tell plotpanel to redraw when idle
        self.plotpanel._resizeflag = True
    
    def draw(self):
        '''Low-level function, draw all graphs from self.graphs. Parameters
        have been set in self.Redraw.'''
        self.figure.clear()
        
        if not self.active:
            self.canvas.draw()
            return
        
        self.subplots = []
        # Draw
        for i, (function, args, kwargs) in enumerate(self.graphs):
            subplot = self.figure.add_subplot(self.rows, self.cols, i + 1)
            function(subplot, *args, **kwargs)
            
            if i < len(self.graphs) - self.cols:
                # X date labels only on lowest graphs
                subplot.xaxis.set_major_formatter(mpldates.DateFormatter(""))
                subplot.xaxis.set_minor_formatter(mpldates.DateFormatter(""))
                subplot.set_xlabel("")

            if i % self.cols != 0:
                # Y label only in left column
                subplot.set_ylabel("")
            
            self.subplots.append(subplot)
        
        self.canvas.draw()

class MultiDataGraph(MultiGraph):
    '''MultiGraph with data collection'''
    def clear(self):
        self.figure.clear()
        self.canvas.draw()
    
    def plot_group(self, groupid, fromdate, todate):
        '''Set plot parameters for group plot'''
        animals = animals_in_group(groupid)
        
        yMax  = animals_getmaxfeeds(animals, fromdate, todate)
        yMax = yMax / 1000. + 0.5
        
        self.graphs = []
        for animalid in animals:
            self.graphs.append((plot_animal, [animalid, fromdate, todate],
                                {'yMax': yMax}))
        
        self.Redraw()

    def plot_section(self, section, fromdate, todate):
        '''Set plot parameters for section plot'''
        self.graphs = []
        
        # Plot per-animal feeds for groups with animals
        afeedgroups = [groupid for groupid in groups_in_section(section)
                       if animals_in_group(groupid)]
        
        for groupid in afeedgroups:
            self.graphs.append((plot_group, [groupid, fromdate, todate], {}))
        
        # Plot groupfeed data for pens without per-animal data
        gfeedMax = 0
        for valve, vgroups in pens_per_valve.items():
            groupids = [section * 10 + pen for pen in vgroups]
            
            if [g for g in groupids if g in afeedgroups]:
                # Already in per-animal data
                continue
            
            groupid = section * 10 + valve
            
            data = groupfeed_average([groupid], fromdate, todate)
            if not data:
                # No groupfeed data
                continue
            
            dMax = max([g for d, g in data])
            if dMax > gfeedMax:
                gfeedMax = dMax
            
            self.graphs.append((plot_groupfeed, [groupid, fromdate, todate],{}))
        
        # Fill in yMax
        animals = animals_in_section(section)
        afeedMax = animals_getmaxfeeds(animals, fromdate, todate)
        yMax = max(afeedMax, gfeedMax) / 1000. + 0.5
        
        for graph in self.graphs:
            graph[2]['yMax'] = yMax
        
        # Sort by groupid
        self.graphs.sort(key = lambda f: f[1][0])
        
        self.Redraw()

    def plot_all(self, fromdate, todate):
        '''Set plot parameters for multi-section plot'''
        sections = sectionlist()
        
        yMax = animals_getmaxfeeds(None, fromdate, todate)
        yMax = yMax / 1000. + 0.5
        
        self.graphs = []
        for section in sections:
            self.graphs.append((plot_section, [section, fromdate, todate],
                                {'yMax': yMax}))
        
        self.Redraw()


if __name__ == '__main__':
    
    class MyApp(wx.App):
        def OnInit(self):
            wx.InitAllImageHandlers()
            
            todate = datetime.date(2007,07,27)
            fromdate = todate - datetime.timedelta(14)
            
            frame = wx.Frame(None, -1, "GroupGraph unit testing")
            frame.SetSize((800, 700))
            d = MultiDataGraph(frame)
            d.plot_group(136, fromdate, todate)
            frame.Show(True)
            self.SetTopWindow(frame)
            
            return True

    app = MyApp(0)
    app.MainLoop()