#!/usr/bin/env python

### FIXME
### replace comboboxes with file menu

import re,sys,os
from optparse import OptionParser
import matplotlib
matplotlib.use("GTK")
from matplotlib.figure import Figure 
from matplotlib.axes import Subplot 
from matplotlib.backends.backend_gtk import FigureCanvasGTK, NavigationToolbar 
from numpy import arange,sin, pi 
import pygtk 
pygtk.require("2.0") 
import gtk 
import gtk.glade
import gobject
from pylab import *
import gnome
from matplotlib import patches
import scipy

import ocrolib

from matplotlib.backends.backend_gtk import FigureCanvasGTK as FigureCanvas
from matplotlib.backends.backend_gtk import NavigationToolbar2GTK as NavigationToolbar

#from matplotlib.backends.backend_gtkcairo import FigureCanvasGTKCairo as FigureCanvas
#from matplotlib.backends.backend_gtkcairo import NavigationToolbar2Cairo as NavigationToolbar
#from matplotlib.backends.backend_gtkagg import FigureCanvasGTKAgg as FigureCanvas
#from matplotlib.backends.backend_gtkagg import NavigationToolbar2GTKAgg as NavigationToolbar

parser = OptionParser(usage="""
%prog [options] [input.db]

Trains models based on a cluster database.
""")
parser.add_option("-v","--verbose",help="verbose",action="store_true")
(options,args) = parser.parse_args()

if len(args)<1:
    parser.print_help()
    sys.exit(0)

iconwidth = 200

def gtk_yield():
    while gtk.events_pending():
       gtk.main_iteration(False)

def numpy2pixbuf(a):
    """Convert a numpy array to a pixbuf."""
    if len(a.shape)==3:
        data = zeros(list(a.shape),'B')
        data[:,:,:] = 255*a
        return gtk.gdk.pixbuf_new_from_array(data,gtk.gdk.COLORSPACE_RGB,8)
    elif len(a.shape)==2:
        data = zeros(list(a.shape)+[3],'B')
        data[:,:,0] = 255*a
        data[:,:,1] = 255*a
        data[:,:,2] = 255*a
        return gtk.gdk.pixbuf_new_from_array(data,gtk.gdk.COLORSPACE_RGB,8)

def get_icon(file,size=iconwidth):
    image = ocrolib.read_image_gray(file)
    print mean(image)
    h,w = image.shape[:2]
    scale = max(w/size,h/size)
    scale = max(1.0,scale)
    if scale>1.0:
        icon = scipy.ndimage.interpolation.zoom(image,(1.0/scale,1.0/scale))
        result = numpy2pixbuf(icon)
        del icon
    else:
        result = numpy2pixbuf(image)
    del image
    return result
    
def compute_component_combolist(widget,kind=None,dflt=None):
    """Compute the combolist from the current charlist."""
    combolist = gtk.ListStore(str)
    cl = ocrolib.ComponentList()
    for i in range(cl.length()):
        if cl.kind(i)!=kind: continue
        if cl.name(i)==dflt: select = len(combolist)
        name = " "+cl.name(i); name = name[1:]
        combolist.append([name])
    widget.set_model(combolist)
    widget.set_text_column(0)
    widget.set_active(select)

class GraphWindow: 
    def __init__(self): 
        self.file = None
        gladefile = ocrolib.findfile("ocropus-showpsegs.glade")
        self.windowname = "pagesegs" 
        self.wtree = gtk.glade.XML(gladefile, self.windowname) 
        self.window = self.wtree.get_widget(self.windowname)
        dic = {
            "on_window1_destroy" : gtk.main_quit,
            "on_pagelist_item_activated" : self.pagelistActivated,
            "on_segselect_changed" : self.segupdate,
            "on_ppselect_changed" : self.ppupdate,
            }
        self.pagelist = self.wtree.get_widget("pagelist")
        self.pagelist.set_selection_mode(gtk.SELECTION_MULTIPLE)
        self.pagelist.set_item_width(iconwidth)
        self.pagelist.set_pixbuf_column(0)
        self.pagelist.set_text_column(1)
        self.pagelist.show_all()

        self.ppselect = self.wtree.get_widget("ppselect")
        compute_component_combolist(self.ppselect,"IBinarize","StandardPreprocessing")
        self.ppupdate()
        self.segselect = self.wtree.get_widget("segselect")
        compute_component_combolist(self.segselect,"ISegmentPage","SegmentPageByRAST")
        self.segupdate()

        self.notebook = self.wtree.get_widget("notebook")
        self.wtree.signal_autoconnect(dic)
        self.window.show_all()
    def removeChildren(self,viewport):
        if type(viewport) is str:
            viewport = self.wtree.get_widget(viewport)
        if not viewport: return
        for child in viewport.get_children():
            viewport.remove(child)
    def removeFigures(self):
        for x in ["raw","cleanup","segmentation"]:
            self.removeChildren(x+"_container")
            self.removeChildren(x+"_toolbar")
    def makeFigure(self,where):
        # setup matplotlib stuff on first notebook page (empty graph) 
        # figure = Figure(figsize=(6,4), dpi=72)
        figure = Figure()
        axis = figure.add_subplot(111)
        canvas = FigureCanvas(figure) 
        canvas.show()         
        viewport = self.wtree.get_widget(where)
        viewport.add(canvas)
        toolname = re.sub("_container","_toolbar",where)
        toolport = self.wtree.get_widget(toolname)
        if toolport is not None:
            toolbar = NavigationToolbar(canvas,self.window)
            self.removeChildren(toolname)
            toolport.add(toolbar)
        return axis
    def addImages(self,images):
        """Set the store for the target class."""
        self.grid = gtk.ListStore(gtk.gdk.Pixbuf,
                                  str,
                                  gobject.TYPE_PYOBJECT)
        self.pagelist.set_model(self.grid)
        for image in images:
            icon = get_icon(image)
            row = [icon,image,None]
            self.grid.append(row)
            gtk_yield()
    def pagelistActivated(self,view,index):
        file = self.grid[index][1]
        self.file = file
        self.recompute()
    def ppupdate(self,*args):
        name = self.ppselect.get_active_text()
        self.preproc = ocrolib.make_IBinarize(name)
        self.recompute()
    def segupdate(self,*args):
        name = self.segselect.get_active_text()
        self.segmenter = ocrolib.make_ISegmentPage(name)
        self.recompute()
    def recompute(self):
        file = self.file
        self.removeFigures()
        if file is None: return
        print self.file,self.preproc.name(),self.segmenter.name()
        
        # load the image
        print "loading",file
        page_gray = ocrolib.read_image_gray(file)
        self.raw_axis = self.makeFigure("raw_container")
        self.raw_axis.imshow(page_gray,cmap=cm.gray)
        gtk_yield()
        
        # clean it up using the given cleanup method
        print "cleanup",file
        page_bin,page_gray = self.preproc.binarize(page_gray)
        self.cleanup_axis = self.makeFigure("cleanup_container")
        self.cleanup_axis.imshow(page_bin,cmap=cm.gray)
        gtk_yield()
        
        # perform segmentation
        print "segmentation",file
        self.segmentation_axis = self.makeFigure("segmentation_container")
        self.segmentation_axis.imshow(page_bin,cmap=cm.gray)
        page_seg = self.segmenter.segment(page_bin)      
        regions = ocrolib.RegionExtractor()
        regions.setPageLines(page_seg)
        print "lines:",regions.length()
        h,w = page_bin.shape
        for i in range(1,regions.length()):
            x0,y0,x1,y1 = regions.bboxMath(i)
            p = patches.Rectangle((x0,h-y1-1),x1-x0,y1-y0,edgecolor="red",fill=0)
            self.segmentation_axis.add_patch(p)
        gtk_yield()
        page_gray.clear()
        page_bin.clear()
        page_seg.clear()
        print "done"

def main():
    app = GraphWindow()
    app.addImages(args)
    gtk.main()

main()
