#!/usr/bin/env python
################################################################################
#
#       This file is part of ghmm
#
#       file:   HMMEd.py
#       author: Janne Grunau
#
#       Copyright (C) 1998-2002, Alexander Schliep
#                                   
#       Contact: schliep@molgen.mpg.de
#
#       Information: http://ghmm.org
#
#       This library is free software; you can redistribute it and/or
#       modify it under the terms of the GNU Library General Public
#       License as published by the Free Software Foundation; either
#       version 2 of the License, or (at your option) any later version.
#
#       This library 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
#       Library General Public License for more details.
#
#       You should have received a copy of the GNU Library General Public
#       License along with this library; if not, write to the Free
#       Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
#
#
#
#       This file is version $Revision: 1.29 $ 
#                       from $Date: 2005/02/22 11:12:56 $
#             last change by $Author: schliep $.
#
################################################################################

import Tkinter
import math
from Gato import GatoGlobals, ProbEditorBasics, ProbEditorDialogs, ProbEditorWidgets, ProbEditorContinuous
from Gato.ObjectGred import *
import ObjectHMM

import ghmmwrapper


class ContinuousEmissionEditor(ProbEditorContinuous.gauss_editor):

    def __init__(self, master, plotList):
        """
        the entire dialog
        """
        ProbEditorContinuous.gauss_editor.__init__(self, master, plotList, width=300, height=300)
        self.normalize = Tkinter.IntVar()
        self.normalize.set(1)


    def buildMenu(self):        
        #Menuleiste
        bar=Tkinter.Menu(self.root)

        filem=Tkinter.Menu(bar)
        editm=Tkinter.Menu(bar)
        zoomm=Tkinter.Menu(bar)

        addMenu=Tkinter.Menu(editm)
        delMenu=Tkinter.Menu(editm)
        if self.sumindi==0:
            addMenu.add_radiobutton(label="Sum-Fkt", command=self.add_sum)

        addMenu.add_radiobutton(label="Uniform", command=self.boxadd)
        #addMenu.add_radiobutton(label="Exp-Fkt", command=self.expadd)
        #addMenu.add_radiobutton(label="NegExp-Fkt", command=self.oexpadd)
        addMenu.add_radiobutton(label="Gaussian", command=self.gaussadd)
        addMenu.add_radiobutton(label="Gaussian left tail", command=self.gaussladd)
        addMenu.add_radiobutton(label="Gaussian right tail", command=self.gaussradd)
        for i in xrange(len(self.plot_list)):
            delMenu.add_radiobutton(label=str(i+1), background = self.plot_list[i].color, command = self.make_del_function(i))
        if self.sumindi==1:
            delMenu.add_radiobutton(label="sum", background='red', command=self.del_sum)

        filem.add_command(label="Save", command=self.save)
        filem.add_command(label="Cancel", command=self.cancel)

        editm.add_cascade(label="Add", menu=addMenu)
        editm.add_cascade(label="Del", menu=delMenu)

        zoomm.add_command(label="Zoom-in", command=self.zoom_in)
        zoomm.add_command(label="Zoom-out", command=self.zoom_out)
        zoomm.add_command(label="Normalise", command=self.norm)

        bar.add_cascade(label="File", menu=filem)
        bar.add_cascade(label="Edit", menu=editm)
        bar.add_cascade(label="Zoom", menu=zoomm)

        self.root.config(menu=bar)

    def boxadd(self):
        self.top=Tkinter.Toplevel(self.root)
        label=Tkinter.Frame(self.top)

        Tkinter.Label(label, justify=CENTER, text="Uniform density:\nf(x) = 1/|start-end| for start<x<end\nelse f(x) = 0").grid(row=0, columnspan=2)
        Tkinter.Label(label, text="start = ").grid(row=1, sticky=E)
        Tkinter.Label(label, text="end = ").grid(row=2, sticky=E)
        self.e1=Tkinter.Entry(label)
        self.e2=Tkinter.Entry(label)
        self.e1.insert(0, -0.5)
        self.e2.insert(0, 0.5)
        self.e1.grid(row=1, column=1)
        self.e2.grid(row=2, column=1)

        button1=Tkinter.Button(label, text="OK", command=self.boxc).grid(row=3)
        button2=Tkinter.Button(label, text="Cancel", command=self.top.destroy).grid(row=3, column=1)

        label.pack()

    def boxc(self):
        s1 = string.atof(self.e1.get())
        s2 = string.atof(self.e2.get())
        self.create_new_fkt(
            ObjectHMM.UniformDensity(start=s1,stop=s2,a=1/math.fabs(s2-s1),
                                     color=self.nextColor()))
        self.top.destroy()

    def gaussadd(self):
        self.top=Tkinter.Toplevel(self.root)
        label=Tkinter.Frame(self.top) 
        
        Tkinter.Label(label, justify=CENTER, text="Normal density:\n f(x) = 1/(sigma*sqrt(2*pi))*exp(-(x-mu)**2/2*(sigma)**2)").grid(row=0, columnspan=2)  
        Tkinter.Label(label, text="mu = ").grid(row=1, sticky=E)
        Tkinter.Label(label, text="sigma = ").grid(row=2, sticky=E)
        Tkinter.Label(label, text="").grid(row=3, sticky=E)
        self.e1=Tkinter.Entry(label)
        self.e2=Tkinter.Entry(label)
        self.e1.insert(0, 0.0)
        self.e2.insert(0, 2.0)
        self.e1.grid(row=1, column=1)
        self.e2.grid(row=2, column=1)

        button1=Tkinter.Button(label, text="OK", command=self.gauss).grid(row=4)
        button2=Tkinter.Button(label, text="Cancel", command=self.top.destroy).grid(row=4, column=1)

        self.label = label
        label.pack()

    def gauss(self):
        mu    = string.atof(self.e1.get())
        sigma = string.atof(self.e2.get())

        if sigma <= 0:
            Tkinter.Label(self.label, text="Error! sigma has to be greater than 0,",
                          foreground="red").grid(row=3, columnspan=2)
            self.label.pack()
            return

        self.create_new_fkt(
            ObjectHMM.NormalDensity(mu=mu,sigma=sigma,a=1.0,color=self.nextColor()))

        del self.label
        self.top.destroy()
        
        
    def gaussladd(self):
        self.top=Tkinter.Toplevel(self.root)
        label=Tkinter.Frame(self.top)
        
        Tkinter.Label(label, justify=CENTER, text="Normal density (left tail):\n f(x) = 1/(sigma*sqrt(2*pi))*exp(-(x-mu)**2/2*(sigma)**2) for x<=tail\n else f(x)=0").grid(row=0, columnspan=2, sticky=E)  
        Tkinter.Label(label, text="mu = ").grid(row=1, sticky=E)
        Tkinter.Label(label, text="sigma = ").grid(row=2, sticky=E)
        Tkinter.Label(label, text="tail = ").grid(row=3, sticky=E)
        Tkinter.Label(label, text="").grid(row=4, sticky=E)
        self.e1=Tkinter.Entry(label)
        self.e2=Tkinter.Entry(label)
        self.e3=Tkinter.Entry(label)
        self.e1.insert(0, -1.0)
        self.e2.insert(0, 1.0)
        self.e3.insert(0, 0.5)
        self.e1.grid(row=1, column=1)
        self.e2.grid(row=2, column=1)
        self.e3.grid(row=3, column=1)

        button1=Tkinter.Button(label, text="OK", command=self.gaussl).grid(row=5)
        button2=Tkinter.Button(label,text="Cancel",command=self.top.destroy).grid(row=5, column=1)

        self.label = label
        label.pack()

    def gaussl(self):
        mu    = string.atof(self.e1.get())
        sigma = string.atof(self.e2.get())
        tail  = string.atof(self.e3.get())

        if sigma <= 0:
            Tkinter.Label(self.label, text="Error! sigma has to be greater than 0,",
                          foreground="red").grid(row=4, columnspan=2)
            self.label.pack()
            return

        self.create_new_fkt(
            ObjectHMM.NormalDensityTruncLeft(mu=mu,sigma=sigma,tail=tail,a=1.0,
                                             color=self.nextColor()))

        del self.label
        self.top.destroy()

    def gaussradd(self):
        self.top=Tkinter.Toplevel(self.root)
        label=Tkinter.Frame(self.top)
        Tkinter.Label(label, justify=CENTER, text="Normal density (right tail):\n f(x) = 1/(sigma*sqrt(2*pi))*exp(-(x-mu)**2/2*(sigma)**2) for x>=tail\n else f(x)=0").grid(row=0, columnspan=2, sticky=E)  
        Tkinter.Label(label, text="mu = ").grid(row=1, sticky=E)
        Tkinter.Label(label, text="sigma = ").grid(row=2, sticky=E)
        Tkinter.Label(label, text="tail = ").grid(row=3, sticky=E)
        Tkinter.Label(label, text="").grid(row=4, sticky=E)
        self.e1=Tkinter.Entry(label)
        self.e2=Tkinter.Entry(label)
        self.e3=Tkinter.Entry(label)
        self.e1.insert(0, 1.0)
        self.e2.insert(0, 1.0)
        self.e3.insert(0, -0.5)
        self.e1.grid(row=1, column=1)
        self.e2.grid(row=2, column=1)
        self.e3.grid(row=3, column=1)
        button1=Tkinter.Button(label, text="OK", command=self.gaussr).grid(row=5)
        button2=Tkinter.Button(label,text="Cancel",command=self.top.destroy).grid(row=5, column=1)  

        self.label = label
        label.pack()

    def gaussr(self):
        mu    = string.atof(self.e1.get())
        sigma = string.atof(self.e2.get())
        tail  = string.atof(self.e3.get())
        
        if sigma <= 0:
            Tkinter.Label(self.label, text="Error! sigma has to be greater than 0,",
                          foreground="red").grid(row=4, columnspan=2)
            self.label.pack()
            return

        self.create_new_fkt(
            ObjectHMM.NormalDensityTruncRight(mu=mu,sigma=sigma,tail=tail,a=1.0,
                                              color=self.nextColor()))
        del self.label
        self.top.destroy()

    def success(self):
        return self.status    


class HMMTypeChooser(tkSimpleDialog.Dialog):
    
    def buttonbox(self):
        # Stolen from tkSimpleDialog.py
        # add standard button box. override if you don't want the
        # standard buttons
        box = Frame(self)
        b_open = Button(box, text="Open existing HMM", width=30, command=self.open, default=ACTIVE)
        b_open.pack(side=TOP, padx=5, pady=5)
        b_discrete = Button(box, text="New HMM with discrete emissions", width=30, command=self.discrete)
        b_discrete.pack(side=TOP, padx=5, pady=5)
        b_continuous = Button(box, text="New HMM with continuous emissions", width=30, command=self.continuous)
        b_continuous.pack(side=BOTTOM, padx=5, pady=5)
        box.pack(side=BOTTOM,fill=X)

    def open(self):
        self.etype = 'open'
        self.ok()

    def discrete(self):
        self.etype = 0
        self.ok()

    def continuous(self):
        self.etype = 1
        self.ok()

    def body(self, master):
        self.title("New HMM")


class HMMGraphEditor(SAGraphEditor):

    def __init__(self, master=None):
        self.modeltype = 0;
        SAGraphEditor.__init__(self, master)

        GatoGlobals.AnimationParameters.cEdgeDefault   = '#888888'
        GatoGlobals.AnimationParameters.cLabelDefault  = 'black' #'#FF8000' # 230 215 0
        GatoGlobals.AnimationParameters.cVertexDefault = 'red' #'#007EE0' # 0 125 230

        self.SetTitle("HMMEd _VERSION_")
        
        self.G = ObjectHMM.ObjectHMM(ObjectHMM.State, ObjectHMM.Transition)

    def makeMenuBar(self):
        self.menubar = Menu(self,tearoff=0)

        # Add file menu
        self.fileMenu = Menu(self.menubar, tearoff=0)
        self.fileMenu.add_command(label='New',            command=self.NewGraph)
        self.fileMenu.add_command(label='Open ...',       command=self.OpenGraph)
        self.fileMenu.add_command(label='Save',	          command=self.SaveGraph)
        self.fileMenu.add_command(label='Save as ...',    command=self.SaveAsGraph)
        self.fileMenu.add_separator()
        self.fileMenu.add_command(label='Export EPSF...', command=self.ExportEPSF)
        self.fileMenu.add_separator()
        self.fileMenu.add_command(label='Quit',	          command=self.Quit)
        self.menubar.add_cascade(label="File", menu=self.fileMenu, underline=0)

        # Add graph menu
        self.graphMenu = Menu(self.menubar, tearoff=0)
        self.graphMenu.add_command(label='Edit Prior', command=self.EditPrior)
        if self.modeltype & ghmmwrapper.kDiscreteHMM:
            self.graphMenu.add_command(label='Edit Alphabet', command=self.EditAlphabet)
        if self.modeltype & ghmmwrapper.kLabeledStates:
            self.graphMenu.add_command(label='Edit class labels', command=self.EditClassLabels)
        if self.modeltype & ghmmwrapper.kBackgroundDistributions:
            self.graphMenu.add_command(label='Edit background distributions', command=self.EditBackgroundDistributions)
        if self.modeltype & ghmmwrapper.kTiedEmissions:
            self.graphMenu.add_command(label='Edit tie groups', command=self.EditTieGroups)
        self.graphMenu.add_separator()
        self.graphMenu.add_checkbutton(label='Grid', command=self.ToggleGridding)
        self.menubar.add_cascade(label="HMM", menu=self.graphMenu, underline=0)

        self.master.configure(menu=self.menubar)

    def SetGraphMenuOptions(self):
        if not self.gridding:
            self.graphMenu.invoke(self.graphMenu.index('Grid'))	


    def CreateWidgets(self):
        toolbar = Frame(self, cursor='hand2', relief=FLAT)
        toolbar.pack(side=LEFT, fill=Y) # Allows horizontal growth

        extra = Frame(toolbar, cursor='hand2', relief=SUNKEN, borderwidth=2)
        extra.pack(side=TOP) # Allows horizontal growth
        extra.rowconfigure(5,weight=1)
        extra.bind("<Enter>", lambda e, gd=self:gd.DefaultInfo())

        px = 0 
        py = 3 

        self.toolVar = StringVar()
        self.lastTool = None

        # Load Icons
        # 0 = "inactive", 1 = "mouse over", 2 = "active"
        self.icons = {
            'AddOrMoveVertex':[PhotoImage(data=GatoIcons.vertex_1),
                               PhotoImage(data=GatoIcons.vertex_2),
                               PhotoImage(data=GatoIcons.vertex_3)],
            'AddEdge':[PhotoImage(data=GatoIcons.edge_1),
                       PhotoImage(data=GatoIcons.edge_2),
                       PhotoImage(data=GatoIcons.edge_3)],
            'DeleteEdgeOrVertex':[PhotoImage(data=GatoIcons.delete_1),
                                  PhotoImage(data=GatoIcons.delete_2),
                                  PhotoImage(data=GatoIcons.delete_3)],
            'SwapOrientation':[PhotoImage(data=GatoIcons.swap_1),
                               PhotoImage(data=GatoIcons.swap_2),
                               PhotoImage(data=GatoIcons.swap_3)],
            'EditWeight':[PhotoImage(data=GatoIcons.edit_1),
                          PhotoImage(data=GatoIcons.edit_2),
                          PhotoImage(data=GatoIcons.edit_3)],
            'EditProperties':[PhotoImage(data=GatoIcons.edit_1),
                              PhotoImage(data=GatoIcons.edit_2),
                              PhotoImage(data=GatoIcons.edit_3)] }
        self.buttons = {}
        values = ['AddOrMoveVertex','AddEdge','DeleteEdgeOrVertex',
                  'SwapOrientation','EditWeight', 'EditProperties']

        text = {'AddOrMoveVertex':'Add or move vertex','AddEdge':'Add edge',
                'DeleteEdgeOrVertex':'Delete edge or vertex',
                'SwapOrientation':'Swap orientation','EditWeight':'Edit Weight',
                'EditProperties':'Edit Properties'}

        row = 0
        for val in values:
            b = Radiobutton(extra, width=32, padx=px, pady=py, 
                            text=text[val],  
                            command=self.ChangeTool,
                            var = self.toolVar, value=val, 
                            indicator=0, image=self.icons[val][0],
                            selectcolor="#AFAFAF",)
            b.grid(row=row, column=0, padx=2, pady=2)
            self.buttons[val] = b
            b.bind("<Enter>", lambda e,gd=self:gd.EnterButtonCallback(e))
            b.bind("<Leave>", lambda e,gd=self:gd.LeaveButtonCallback(e))
            row += 1
            
        self.defaultButton = self.buttons['AddOrMoveVertex']
        # default doesnt work as config option
        GraphEditor.CreateWidgets(self)


    def MouseUp(self,event):
        if self.mode == 'AddOrMoveVertex':
            self.AddOrMoveVertexUp(event)
        elif self.mode == 'AddEdge':
            self.AddEdgeUp(event)
        elif self.mode == 'DeleteEdgeOrVertex':
            self.DeleteEdgeOrVertexUp(event)
        elif self.mode == 'SwapOrientation':
            self.SwapOrientationUp(event)
        elif self.mode == 'EditWeight':
            self.EditWeightUp(event)
        elif self.mode == 'EditProperties':
            self.EditPropertiesUp(event)


    def NewGraph(self):
        d = HMMTypeChooser(self.master)

        if d.etype == 'open':
            self.OpenGraph()
            return
        else:
            self.G = ObjectHMM.ObjectHMM(ObjectHMM.State, ObjectHMM.Transition, etype=d.etype)

        self.G.edit(self.master)

        self.graphName = "New"
        self.ShowGraph(self.G, self.graphName)
        #self.RegisterGraphInformer(HMMInformer(self.G))
        self.fileName = None
        self.SetTitle("HMMEd _VERSION_")
        self.modeltype = self.G.modelType
        self.makeMenuBar()
        self.SetGraphMenuOptions()

    def OpenGraph(self):
        file = askopenfilename(title="Open HMM",
                               defaultextension=".xml",
                               filetypes = (("XML", ".xml"),
                                            )
                               )
        if file is "":
            print "cancelled"
        else:
            self.fileName = file
            self.graphName = stripPath(file)
            e = extension(file)
            
            if e == 'xml':
                self.G.openXML(file)
            else:
                print "Unknown extension"
                return

            self.ShowGraph(self.G, self.graphName)
            #self.RegisterGraphInformer(HMMInformer(self.HMM))
            self.SetTitle("HMMEd _VERSION_ - " + self.graphName)
            self.modeltype = self.G.modelType
            self.makeMenuBar()
            
            if not self.gridding:
                self.graphMenu.invoke(self.graphMenu.index('Grid'))

    def SaveGraph(self):
        #self.dirty = 0
        if self.fileName != None:
            self.G.writeXML(self.fileName)
        else:
            self.SaveAsGraph()
            
    def SaveAsGraph(self):
        file = asksaveasfilename(title="Save HMM",
                                 defaultextension=".xml",
                                 filetypes = ( ("XML", ".xml"),
                                               )
                                 )
        if file is not "":
            self.fileName = file
            self.dirty = 0
            self.G.writeXML(self.fileName)
            self.graphName = stripPath(file)
            self.SetTitle("HMMEd _VERSION_ - " + self.graphName)


    def EditAlphabet(self):
        if self.G.modelType & ghmmwrapper.kHigherOrderEmissions and self.G.Order() > 0:
            print "not implemented, edit the alphabet before adding states"
        else:
            self.G.alphabet.editDialog(self, self.G)


    def EditPrior(self):
        if self.G.Order() == 0:
            return

        emission_probabilities = ProbEditorBasics.ProbDict({})
        for state in self.G.vertices.values():
            label  = str(state.id)
            weight = state.initial
            emission_probabilities.update({label:weight})

        emission_probabilities.renorm_to(1.0)

        e = ProbEditorBasics.emission_data(emission_probabilities)
        d = ProbEditorDialogs.emission_dialog(self, e, "initial probabilities")

        if d.success():
            # write back normalized probabilities
            emission_probabilities.renorm_to(1.0)
            for key in emission_probabilities.keys():
                state = self.G.vertices[int(key)]
                state.initial = typed_assign(state.initial, emission_probabilities[key])


    def EditClassLabels(self):
        self.G.label_alphabet.editDialog(self, self.G)


    def EditBackgroundDistributions(self):
        self.G.backgroundDistributions.editDialog(self, self.G)
        
    def EditTieGroups(self):
        self.G.tie_groups.editDialog(self, self.G)


    def EditWeightUp(self,event):
        if event.widget.find_withtag(CURRENT):
            widget = event.widget.find_withtag(CURRENT)[0]
            tags = self.canvas.gettags(widget)
            if "edges" in tags:
                (tail,head) = self.edge[widget]
                
                transition_probabilities=ProbEditorBasics.ProbDict({})
                for head in self.G.OutNeighbors(tail):
                    weight=self.G.GetEdgeWeight(0,tail,head)
                    label = "-> %d" % head
                    transition_probabilities.update({label:weight})
                    
                if transition_probabilities.sum==0:
                    key_list=transition_probabilities.keys()
                    for key in key_list:
                        transition_probabilities[key]=1.0/len(key_list)
                e = ProbEditorBasics.emission_data(transition_probabilities)
                d = ProbEditorDialogs.emission_dialog(self, e,
                                                      "transition probs from state %d" % tail)
                if d.success():
                    # write back normalized probabilities
                    for key in transition_probabilities.keys():
                        head = int(key[3:])
                        self.G.SetEdgeWeight(0, tail, head, transition_probabilities[key]/transition_probabilities.sum)

            else: # We have a vertex
                v = self.FindVertex(event)
                if v != None:
                    self.G.vertices[v].editEmissions(self.master)

    def EditPropertiesUp(self, event):
        if event.widget.find_withtag(CURRENT):
            widget = event.widget.find_withtag(CURRENT)[0]
            tags = self.canvas.gettags(widget)        
            if "edges" not in tags:
                v = self.FindVertex(event)
                if v != None:
                    self.G.vertices[v].editProperties(self.master)
                    self.UpdateVertexLabel(v)

    def AddVertex(self):
        """ Add an isolated vertex. Returns the id of the new vertex """
        id = ObjectGraph.AddVertex(self)
        self.G.edges[id].itsHMM = self.G
        return id

class Start:

    def __init__(self):
        graphEditor = SAGraphEditorToplevel()
        graphEditor.NewGraph()
        import logging
        log = logging.getLogger("HMMEditor.py")
        
        ################################################################################
if __name__ == '__main__':
    
    # Overide default colors for widgets ... maybe shouldnt be doing that for Windows?
    tk = Tk()
    tk.option_add('*ActiveBackground','#EEEEEE')
    tk.option_add('*background','#DDDDDD')
    tk.option_add('Tk*Scrollbar.troughColor','#CACACA')
    HMMEditor = HMMGraphEditor(tk)
    HMMEditor.NewGraph()
    import logging
    log = logging.getLogger("HMMEditor.py")
    HMMEditor.mainloop()
    
