# GVB - a GTK+/GNOME vibrations simulator
#
# Copyright © 2008-2013 Pietro Battiston
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA


from gi.repository import Gtk, Gdk
from scipy import shape
from time import time
from cairo import ImageSurface, FORMAT_ARGB32, Context

def one_basic(cr, size, pos, speeds = None):
    spaced=1./(len(pos)+1)

    cr.scale(size[0], size[1]/4)
    cr.translate(0, 2)
    cr.set_line_width(.01)
    cr.set_source_rgb(1,0,0)

    cr.move_to(0,0)

    xposition=0
    for ypos in pos:
        xposition=xposition+spaced

        cr.line_to(xposition, -ypos)
    cr.line_to(1,0)
    cr.stroke()


def with_speeds(cr, size, pos, speeds):
    n = len(pos)
    spaced=1./(n+1)
    coeff = .1*n
    cr.scale(size[0], size[1]/4)
    cr.translate(0, 2)
    cr.set_line_width(.1/n)
    cr.set_source_rgb(1,0,0)


    cr.move_to(0,0)

    xposition=spaced
    for idx in range(len(pos)):
        cr.move_to(xposition,-pos[idx])
        cr.rel_line_to(0,-speeds[idx]*coeff)
        cr.stroke()
        xposition=xposition+spaced

    cr.set_line_width(.02)
    cr.set_source_rgb(0,0,1)
    cr.move_to(0,0)

    xposition=0
    for ypos in pos:
        xposition=xposition+spaced

        cr.line_to(xposition, -ypos)
    cr.line_to(1,0)
    cr.stroke()


def temperature(cr, size, pos, speeds = None):
    spaced=1./(len(pos))

    cr.scale(size[0], size[1]/4)
    cr.translate(0, 2)
    cr.set_line_width(.5)
    cr.set_source_rgb(1,0,0)

    cr.move_to(0,0)

    xposition=0
    for ypos in pos:
        cr.move_to(xposition,0)
        xposition=xposition+spaced

        if ypos >= 1:
            cr.set_source_rgb(1,1,1)
        elif ypos > .5:
            cr.set_source_rgb((1-ypos)/2,(1-ypos)/2,0)
        elif ypos > 0:
            cr.set_source_rgb(ypos/2, 1-1.5*ypos, 0)
        elif ypos > -.5:
            cr.set_source_rgb(0,1+2*ypos,-ypos)
        else:
            cr.set_source_rgb(0,0,1+ypos)
        cr.line_to(xposition, 0)
        cr.stroke()


def _two_colours(cr, pos, size, speeds = None):
    hspaced=1./shape(pos)[0]
    vspaced=1./shape(pos)[1]

    cr.scale(size[0], size[1])
    yposition=0
    for ypos in pos:
        xposition=0

        for xpos in ypos:
            if xpos > 0:
                cr.set_source_rgb(1-xpos,1,1)
            else:
                cr.set_source_rgb(1,1+xpos,1)
            cr.rectangle(xposition, yposition, hspaced, vspaced)
            cr.fill()
            xposition=xposition+hspaced
        yposition=yposition+vspaced


def colours(cr, size, pos, speeds = None):
    hspaced=1./shape(pos)[0]
    vspaced=1./shape(pos)[1]

    cr.scale(size[0], size[1])
    cr.translate(0, vspaced/2)
    cr.set_line_width(vspaced)
    yposition=0
    for ypos in pos:
        xposition=0
        cr.move_to(0,yposition)

        for xpos in ypos:
            cr.move_to(xposition, yposition)
            xposition=xposition+hspaced

            if xpos >= 1:
                cr.set_source_rgb(1,1,1)
            elif xpos > .5:
                cr.set_source_rgb((1-xpos)/2,(1-xpos)/2,0)
            elif xpos > 0:
                cr.set_source_rgb(xpos/2, 1-1.5*xpos, 0)
            elif xpos > -.5:
                cr.set_source_rgb(0,1+2*xpos,-xpos)
            else:
                cr.set_source_rgb(0,0,1+xpos)

            cr.line_to(xposition, yposition)
            cr.stroke()
        yposition=yposition+vspaced



def trid(cr, size, pos, speeds = None):
    m,n=shape(pos)

    vspaced=1./(3*n)

    yposition=.3

    cr.set_line_join(2)

    pos = pos * .5

    tr = ([[yposition for i in range(n+2)]]
        +[[yposition+(j+1)*vspaced]+[yposition + (j+1)*vspaced-pos[j][i]/2 for i in range(n)]+[yposition+(j+1)*vspaced] for j in range(m)]
        +[[yposition+(m+1)*vspaced for i in range(n+2)]])

    cr.scale(size[0], size[1])

    cr.set_line_width(.001)

    top_margin = .1

    for row_idx in range(m+1):
        margin = top_margin*(1-(row_idx/m))
        marg_under = top_margin*(1-((row_idx+1)/m))
        hspaced=(1-2*margin)/(n+3)
        xposition = margin + hspaced
        hsp_under=(1-2*marg_under)/(n+3)
        xpos_under = marg_under + hsp_under

        for col_idx in range(n+1):
            if tr[row_idx][col_idx+1] > tr[row_idx+1][col_idx+1]:
                cr.set_source_rgba(1,0,0,.8)
            else:
                cr.set_source_rgba(0,1,0,.8)
            cr.move_to(xposition, tr[row_idx][col_idx])
            cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
            cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
            cr.close_path()
            cr.fill()

            if tr[row_idx][col_idx] > tr[row_idx+1][col_idx]:
                cr.set_source_rgba(1,0,0,.8)
            else:
                cr.set_source_rgba(0,1,0,.8)
            cr.move_to(xposition, tr[row_idx][col_idx])
            cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
            cr.line_to(xpos_under, tr[row_idx+1][col_idx])
            cr.close_path()
            cr.fill()

            cr.set_source_rgb(0,0,0)

            cr.move_to(xposition, tr[row_idx][col_idx])
            cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
            cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
            cr.line_to(xpos_under, tr[row_idx+1][col_idx])
            cr.close_path()
            cr.stroke()
            xposition=xposition+hspaced
            xpos_under=xpos_under+hsp_under

def trid_speeds(cr, size, pos, speeds = None):
    m,n=shape(pos)

    vspaced=1./(3*n)

    coeff = .05*max(m,n)

    yposition=.3

    cr.set_line_join(2)

    pos = pos * .5

    tr = ([[yposition for i in range(n+2)]]
        +[[yposition+(j+1)*vspaced]+[yposition + (j+1)*vspaced-pos[j][i]/2 for i in range(n)]+[yposition+(j+1)*vspaced] for j in range(m)]
        +[[yposition+(m+1)*vspaced for i in range(n+2)]])

    cr.scale(size[0], size[1])
    cr.set_line_width(.001)

    top_margin = .1

    for row_idx in range(m+1):
        margin = top_margin*(1-(row_idx/m))
        marg_under = top_margin*(1-((row_idx+1)/m))
        hspaced=(1-2*margin)/(n+3)
        xposition = margin + hspaced
        hsp_under=(1-2*marg_under)/(n+3)
        xpos_under = marg_under + hsp_under
#        print ypos
        for col_idx in range(n+1):
            cr.set_line_width(.001)
            if tr[row_idx][col_idx+1] > tr[row_idx+1][col_idx+1]:
                cr.set_source_rgba(0,0,1,.8)
            else:
                cr.set_source_rgba(0,1,0,.8)
            cr.move_to(xposition, tr[row_idx][col_idx])
            cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
            cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
            cr.close_path()
            cr.fill()

            if tr[row_idx][col_idx] > tr[row_idx+1][col_idx]:
                cr.set_source_rgba(0,0,1,.8)
            else:
                cr.set_source_rgba(0,1,0,.8)
            cr.move_to(xposition, tr[row_idx][col_idx])
            cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
            cr.line_to(xpos_under, tr[row_idx+1][col_idx])
            cr.close_path()
            cr.fill()

            cr.set_source_rgb(0,0,0)

            cr.move_to(xposition, tr[row_idx][col_idx])
            cr.line_to(xposition + hspaced, tr[row_idx][col_idx+1])
            cr.line_to(xpos_under + hsp_under, tr[row_idx+1][col_idx+1])
            cr.line_to(xpos_under, tr[row_idx+1][col_idx])
            cr.close_path()
            cr.stroke()

            if row_idx != 0 and col_idx != 0:
                cr.set_line_width(.1/n)

                cr.set_source_rgba(1,0,0,.5)
                cr.move_to(xposition, tr[row_idx][col_idx])
                cr.rel_line_to(0,-speeds[row_idx-1][col_idx-1]*coeff)
                cr.stroke()

            xposition=xposition+hspaced
            xpos_under=xpos_under+hsp_under

drawers_1d = {
                'wave' :    one_basic,
                'temp' :    temperature,
                'speeds':    with_speeds
            }

drawers_2d = {
                'geo' :        colours,
                '3D' :        trid,
                'speeds':    trid_speeds
            }


drawers_dict={ 1: drawers_1d, 2: drawers_2d }
drawers = {1: ['wave', 'temp', 'speeds'], 2: ['3D', 'geo', 'speeds']}

class Drawer():
    def __init__(self, drawing, combo, dr_type = None):
        self.drawing = drawing
        self.dumpdir = None
        self.combo = combo
        self.active = True
        self.points = None
        self.dim = None
        self.ready = False
        self.drawn = 0
        self.dr_type = dr_type
        self.combo.connect('changed', self.change_drawer)
        if hasattr(self.drawing, 'signals_ids'):
            # This means that the drawing area was already drawn by another
            # drawer. This one is a possessive guy and want her to forget him.
            for signal_id in self.drawing.signals_ids:
                self.drawing.disconnect(signal_id)
        
        # Throw away all his ids! You're mine now!
        self.drawing.signals_ids = []
        self.drawing.signals_ids.append(self.drawing.connect('draw', self.redraw))
        self.size = None

    def draw(self, cr):
        if self.dim != len(self.points.shape):
            raise
        if not self.drawing.get_window():
            return
        if self.dumpdir:
            try:
                cr.set_source_surface(self._last_surface, 0, 0)
            except AttributeError:
                # Still no surface ready
                pass

        before=time()
        if not hasattr(self, 'drawer'):
            return
        size = self.drawing.get_allocated_width(), self.drawing.get_allocated_height()

        cr.rectangle(0, 0, *size)
        cr.set_source_rgb(1, 1, 1)
        cr.fill()
        
        self.drawer(cr, size, self.points.pos, self.points.speeds)
        self.drawn=time()
        self.drawing_time = self.drawn-before

    def dump_draw(self):

        self._last_surface = ImageSurface(FORMAT_ARGB32, self.size[0], self.size[1])

        self.cr = Context(self._last_surface)

        before=time()
        self.drawer(self.cr, self.size, self.points.pos, self.points.speeds)

        self._last_surface.write_to_png('%s/%08d.png' % (self.dumpdir,
                                                         self.dumpindex))
        self.dumpindex += 1
        self.drawn=time()
        self.drawing_time = self.drawn-before
        self.invalidate()
    
    def invalidate(self):
        self.drawing.queue_draw()
    
    def redraw(self, widget, cr):
        self.draw(cr)

    def change_drawer(self, combo):
        self.dr_type = self.combo.get_active_text()
        if self.dr_type in drawers[self.dim]:
            self.drawer=drawers_dict[self.dim][self.dr_type]
            self.drawing.queue_draw()


    def configure(self, points=None, dim=None):
        if not self.active:
            return

        if points:
            self.points = points
            dim = len(shape(points.pos))

        if dim:
            if self.dim != dim:
                self.dim = dim
                self.build_menu()
                self.ready = True
        else:
            return

#        if not (self.points and len(shape(self.points.pos)) == self.dim):    #I won't change dimension if they don't tell me!
#            return

    def build_menu(self):
        combo_model = Gtk.ListStore(str)
        for item in drawers[self.dim]:
            combo_model.append([item])
        self.combo.set_model(combo_model)
        if self.dr_type in drawers[self.dim]:
            self.combo.set_active(drawers[self.dim].index(self.dr_type))
        else:
            self.combo.set_active(0)

    def dump(self, dumpdir):
        print("dumping in", dumpdir)
        self.dumpindex = 0

        if dumpdir:                #This is a "start" command
            self.dumpdir = dumpdir
            self.size = self.drawing.get_allocated_width(), self.drawing.get_allocated_height()
            self.points.changed = self.dump_draw
        else:                        #This is a "stop" command
            self.dumpdir = None
            self.size = None
            self.points.changed = self.invalidate
