"""
@package wxplot.scatter

@brief Scatter plotting using PyPlot

Classes:
 - scatter::ScatterFrame
 - scatter::ScatterToolbar

(C) 2011 by the GRASS Development Team

This program is free software under the GNU General Public License
(>=v2). Read the file COPYING that comes with GRASS for details.

@author Michael Barton, Arizona State University
"""

import sys

import wx
try:
    import wx.lib.plot as plot
except ImportError as e:
    print >> sys.stderr, e

import grass.script as grass

from wxplot.base       import BasePlotFrame, PlotIcons
from gui_core.toolbars import BaseToolbar, BaseIcons
from wxplot.dialogs    import ScatterRasterDialog, PlotStatsFrame
from core.gcmd         import RunCommand, GException, GError, GMessage
from core.utils import _

class ScatterFrame(BasePlotFrame):
    """Mainframe for displaying bivariate scatter plot of two raster maps. Uses wx.lib.plot.
    """
    def __init__(self, parent, id = wx.ID_ANY, style = wx.DEFAULT_FRAME_STYLE, 
                 size = wx.Size(700, 400),
                 rasterList = [], **kwargs):
        BasePlotFrame.__init__(self, parent, size = size, **kwargs)
        
        self.toolbar = ScatterToolbar(parent = self)
        # workaround for http://trac.wxwidgets.org/ticket/13888
        if sys.platform != 'darwin':
            self.SetToolBar(self.toolbar)
        self.SetTitle(_("GRASS Bivariate Scatterplot Tool"))

        #
        # Init variables
        #
        self.rasterList = rasterList
        self.plottype = 'scatter'
        self.ptitle = _('Bivariate Scatterplot')     # title of window
        self.xlabel = _("Raster cell values")           # default X-axis label
        self.ylabel = _("Raster cell values")           # default Y-axis label
        self.maptype = 'raster'                         # default type of scatterplot
        self.scattertype = 'normal' 
        self.bins = 255
        self.colorList = ["blue", "red", "black", "green", "yellow", "magenta", "cyan", \
                    "aqua", "grey", "orange", "brown", "purple", "violet", \
                    "indigo"]
        
        self._initOpts()

        if len(self.rasterList) > 1: # set raster name(s) from layer manager if a map is selected
            self.InitRasterOpts(self.rasterList, 'scatter')
        else:
            self.raster = {}

    def _initOpts(self):
        """Initialize plot options
        """
        self.InitPlotOpts('scatter')            

    def OnCreateScatter(self, event):
        """Main routine for creating a scatterplot. Uses r.stats to
        create a list of cell value pairs. This is passed to
        plot to create a scatterplot.
        """
        self.SetCursor(wx.StockCursor(wx.CURSOR_ARROW))
        self.SetGraphStyle()
        wx.BeginBusyCursor()
        wx.SafeYield()
        self.SetupScatterplot()
        p = self.CreatePlotList()
        if p:
            self.DrawPlot(p)
            wx.EndBusyCursor()
        else:
            wx.EndBusyCursor()
            GMessage(_("Nothing to plot."), parent = self)

    def OnSelectRaster(self, event):
        """Select raster map(s) to profile
        """
        dlg = ScatterRasterDialog(parent = self)
        dlg.CenterOnParent()
        if dlg.ShowModal() == wx.ID_OK:
            self.rasterList = dlg.GetRasterPairs()
            if not self.rasterList:
                GMessage(_("At least 2 raster maps must be specified"), parent = dlg)
                return
            
            # scatterplot or bubbleplot
            # bins for r.stats with float and dcell maps
            self.scattertype, self.bins = dlg.GetSettings()
            self.raster = self.InitRasterPairs(self.rasterList, 'scatter') # dictionary of raster pairs
            
            # plot histogram
            if self.rasterList:
                self.OnCreateScatter(event = None)
        
        dlg.Destroy()
        
    def SetupScatterplot(self):
        """Build data list for ploting each raster
        """

        #
        # initialize title string
        #
        self.ptitle = _('Bivariate Scatterplot of ')        

        #
        # create a datalist for plotting for each raster pair
        #
        if len(self.rasterList) == 0: return  # at least 1 pair of maps needed to plot        
        
        for rpair in self.rasterList:
            self.raster[rpair]['datalist'] = self.CreateDatalist(rpair)
            
            # update title
            self.ptitle += '%s vs %s, ' % (rpair[0].split('@')[0], rpair[1].split('@')[0])

        self.ptitle = self.ptitle.strip(', ')
        
        #
        # set xlabel & ylabel based on raster maps of first pair to be plotted
        #
        self.xlabel = _('Raster <%s> cell values') % rpair[0].split('@')[0]
        self.ylabel = _('Raster <%s> cell values') % rpair[1].split('@')[0]
        
        units = self.raster[self.rasterList[0]][0]['units']
        if units != '':
            self.xlabel += _(': %s') % units

        units = self.raster[self.rasterList[0]][1]['units']
        if units != '':
            self.ylabel += _(': %s') % units
            
    def CreateDatalist(self, rpair):
        """Build a list of cell value, frequency pairs for histogram
            frequency can be in cell counts, percents, or area
        """
        datalist = []
        
        if self.scattertype == 'bubble': 
            freqflag = 'cn'
        else:
            freqflag = 'n'
                
        try:
            ret = RunCommand("r.stats",
                             parent = self,
                             input = '%s,%s' % rpair,
                             flags = freqflag,
                             nsteps = self.bins,
                             sep = ',',
                             quiet = True,
                             read = True)
            
            if not ret:
                return datalist
            
            for line in ret.splitlines():
                rast1, rast2 = line.strip().split(',')
                rast1 = rast1.strip()
                if '-' in rast1:
                    if rast1[0] == '-':
                        rast1 = '-' + rast1.split('-')[1]
                    else:
                        rast1 = rast1.split('-')[0]

                rast2 = rast2.strip()
                if '-' in rast2:
                    if rast2[0] == '-':
                        rast2 = '-' + rast2.split('-')[1]
                    else:
                        rast2 = rast2.split('-')[0]
                
                rast1 = rast1.encode('ascii', 'ignore')
                rast2 = rast2.encode('ascii', 'ignore')
                    
                datalist.append((rast1,rast2))

            return datalist
        except GException as e:
            GError(parent = self,
                   message = e.value)
            return None
        
    def CreatePlotList(self):
        """Make list of elements to plot
        """
        # graph the cell value, frequency pairs for the histogram
        self.plotlist = []
        
        for rpair in self.rasterList:
            if 'datalist' not in self.raster[rpair] or \
                self.raster[rpair]['datalist'] is None:
                continue
            
            if len(self.raster[rpair]['datalist']) > 0:
                col = wx.Colour(self.raster[rpair]['pcolor'][0],
                                self.raster[rpair]['pcolor'][1],
                                self.raster[rpair]['pcolor'][2],
                                255)
                scatterpoints = plot.PolyMarker(self.raster[rpair]['datalist'],
                                                legend = ' ' + self.raster[rpair]['plegend'],
                                                colour = col,size = self.raster[rpair]['psize'],
                                                fillstyle = self.ptfilldict[self.raster[rpair]['pfill']],
                                                marker = self.raster[rpair]['ptype'])

                self.plotlist.append(scatterpoints)
        
        return self.plotlist
    
    def Update(self):
        """Update histogram after changing options
        """
        self.SetGraphStyle()
        p = self.CreatePlotList()
        if p:
            self.DrawPlot(p)
        else:
            GMessage(_("Nothing to plot."), parent = self)
    
    def OnRegression(self, event):
        """Displays regression information in messagebox
        """
        message = []
        title = _('Regression Statistics for Scatterplot(s)')

        for rpair in self.rasterList:
            if isinstance(rpair, tuple) == False: continue
            rast1, rast2 = rpair
            rast1 = rast1.split('@')[0] 
            rast2 = rast2.split('@')[0] 
            ret = grass.parse_command('r.regression.line', 
                                      mapx = rast1, 
                                      mapy = rast2, 
                                      flags = 'g', quiet = True,
                                      parse = (grass.parse_key_val, { 'sep' : '=' }))
            eqtitle = _('Regression equation for raster map <%(rast1)s> vs. <%(rast2)s>:\n\n') % \
                { 'rast1' : rast1,
                  'rast2' : rast2 }
            eq = '   %s = %s + %s(%s)\n\n' % (rast2, ret['a'], ret['b'], rast1)
            num = 'N = %s\n' % ret['N']
            rval = 'R = %s\n' % ret['R']
            rsq = 'R-squared = %f\n' % pow(float(ret['R']), 2)
            ftest = 'F = %s\n' % ret['F']
            str = eqtitle + eq + num + rval + rsq + ftest
            message.append(str)
            
        stats = PlotStatsFrame(self, id = wx.ID_ANY, message = message, 
                                      title = title)

        if stats.Show() == wx.ID_CLOSE:
            stats.Destroy()       

class ScatterToolbar(BaseToolbar):
    """Toolbar for bivariate scatterplots of raster map pairs
    """
    def __init__(self, parent):
        BaseToolbar.__init__(self, parent)

        # workaround for http://trac.wxwidgets.org/ticket/13888
        if sys.platform == 'darwin':
            parent.SetToolBar(self)

        self.InitToolbar(self._toolbarData())
        
        # realize the toolbar
        self.Realize()
        
    def _toolbarData(self):
        """Toolbar data"""
        return self._getToolbarData((('addraster', BaseIcons["addRast"],
                                      self.parent.OnSelectRaster),
                                     (None, ),
                                     ('draw', PlotIcons["draw"],
                                      self.parent.OnCreateScatter),
                                     ('erase', BaseIcons["erase"],
                                      self.parent.OnErase),
                                     ('drag', BaseIcons['pan'],
                                      self.parent.OnDrag),
                                     ('zoom', BaseIcons['zoomIn'],
                                      self.parent.OnZoom),
                                     ('unzoom', BaseIcons['zoomBack'],
                                      self.parent.OnRedraw),
                                     (None, ),
                                     ('statistics', PlotIcons['statistics'],
                                      self.parent.OnRegression),
                                     ('image', BaseIcons["saveFile"],
                                      self.parent.SaveToFile),
                                     ('print', BaseIcons["print"],
                                      self.parent.PrintMenu),
                                     (None, ),
                                     ('settings', PlotIcons["options"],
                                      self.parent.PlotOptionsMenu),
                                     ('quit', PlotIcons["quit"],
                                      self.parent.OnQuit),
                                     ))
