
import numpy
import Tkinter 
import tkMessageBox, tkFileDialog, ScrolledText
import matplotlib
matplotlib.use('TkAgg')
from matplotlib.figure import Figure
import code,os

from pyOM import pyOM

class pyOM_gui(Tkinter.Frame,pyOM):
   """
   pyOM with graphical user interface
   """
   def __init__(self,snapint = 0, plotint = 0):
     pyOM.__init__(self)
     self.pyOM_vstr = 'pyOM v%2.1f' % self.pyOM_version
     Tkinter.Frame.__init__(self,master=None)
     self.pack(expand=Tkinter.YES)
     self.master.title(self.pyOM_vstr)
     self.master.protocol('WM_DELETE_WINDOW', self.quit)

     menubar = Tkinter.Menu(self)
     menu = Tkinter.Menu(menubar, tearoff=0)
     menu.add_command(label="Read restart file", command=self.read)
     menu.add_command(label="Write restart file", command=self.write)
     menu.add_command(label="Save plot to file", command=self.printplot)
     menu.add_separator()
     menu.add_command(label="Exit", command=self.quit)
     menubar.add_cascade(label="File", menu=menu)
     menu = Tkinter.Menu(menubar, tearoff=0)
     menu.add_command(label="Start integration", command=self.run)
     menu.add_command(label="Stop integration", command=self.stop)
     menu.add_command(label="Reset integration", command=self.reset)
     menu.add_command(label="Diagnose", command=self.diagnose)
     menu.add_command(label="Plot", command=self.diag_plot)
     menubar.add_cascade(label="Control", menu=menu)
     menu = Tkinter.Menu(menubar, tearoff=0)
     menu.add_command(label="Shell", command=self.shell)
     menu.add_command(label="Config", command=self.config)
     menubar.add_cascade(label="Options", menu=menu)
     self.master.config(menu=menubar) 
     
     self.image_height=480; self.image_width=640
     self.main_area= Tkinter.Canvas(self,height=self.image_height,width=self.image_width)
     self.main_area.pack(side=Tkinter.TOP)
     self._tkphoto = Tkinter.PhotoImage(master=self.main_area)
     
 
     self.time_step_list = ScrolledText.ScrolledText(self,height=3,state=Tkinter.NORMAL,bd=2)
     self.time_step_list.pack(side=Tkinter.TOP,fill=Tkinter.X,expand=Tkinter.YES)
     
     M=self.fortran.pyom_module         # fortran module with model variables
     self.time_step_list.insert(Tkinter.END,'Welcome to %s\n\n' % self.pyOM_vstr)
     self.time_step_list.insert(Tkinter.END,'grid size    : nx=%i ny=%i nz=%i \n' % (M.nx,M.ny,M.nz)  )
     self.time_step_list.insert(Tkinter.END,'grid spacing : Delta x =%f m  Delta z =%f m\n'   % (M.dx,M.dz)  )
     self.time_step_list.insert(Tkinter.END,'domain size  : %f m X %f m X %f m\n' % (M.nx*M.dx,M.ny*M.dx,M.nz*M.dz) )
     self.time_step_list.insert(Tkinter.END,'time step    : %f s\n' % M.dt )
     self.time_step_list.insert(Tkinter.END,'lateral  diffusivity : K_h= %f m^2/s\n'  % M.k_h  )
     self.time_step_list.insert(Tkinter.END,'vertical diffusivity : K_v= %f m^2/s\n'  % M.k_v  )
     self.time_step_list.insert(Tkinter.END,'lateral  viscosity   : A_h= %f m^2/s\n'  % M.a_h  )
     self.time_step_list.insert(Tkinter.END,'vertical viscosity   : A_v= %f m^2/s\n'  % M.a_v  )
     self.time_step_list.insert(Tkinter.END,'epsilon for 2D solver : %f \n' % M.eps2d_sor   )
     self.time_step_list.yview(Tkinter.MOVETO,1)
     self.time_step_list.config(state=Tkinter.DISABLED)

     if snapint == 0: snapint = 10
     if plotint == 0: plotint = snapint
     self.snapint = Tkinter.IntVar(); self.snapint.set( snapint )
     self.plotint = Tkinter.IntVar(); self.plotint.set( plotint )
     self.halt     = 1
     
     Toolbar = Tkinter.Frame(self,relief=Tkinter.SUNKEN,bd=2)
     Toolbar.pack(side=Tkinter.TOP, expand=Tkinter.YES, fill=Tkinter.X)

     file = os.path.join(matplotlib.rcParams['datapath'], 'images', 'stock_right.ppm')
     self.image_button1 = Tkinter.PhotoImage(master=Toolbar, file=file)
     file = os.path.join(matplotlib.rcParams['datapath'], 'images', 'stock_close.ppm')
     self.image_button2 = Tkinter.PhotoImage(master=Toolbar, file=file)
     
     self.itt_str = Tkinter.StringVar()
     F = Tkinter.Frame(Toolbar); F.pack(side=Tkinter.LEFT)
     Tkinter.Label(F, text = self.pyOM_vstr,width=18).pack(side=Tkinter.TOP,anchor='w')
     Tkinter.Label(F, textvariable = self.itt_str,width=18).pack(side=Tkinter.TOP,anchor='w')
     Tkinter.Button(Toolbar, text = "Run" , image=self.image_button1,command=self.run).pack(side=Tkinter.LEFT)
     Tkinter.Button(Toolbar, text = "Stop", image=self.image_button2,command=self.stop).pack(side=Tkinter.LEFT)
     
     Tkinter.Scale(Toolbar,variable = self.snapint,label="Diagnostic interval", length=120, \
           from_ = 1, to = max(snapint,500), orient=Tkinter.HORIZONTAL).pack(side=Tkinter.LEFT)
     Tkinter.Checkbutton(Toolbar,command=self.scale_both).pack(side=Tkinter.LEFT)
     self.scale_bar = Tkinter.Scale(Toolbar,variable = self.plotint,label="Plotting interval", length=120, \
           from_ = 1, to = max(plotint,500), orient=Tkinter.HORIZONTAL)
     self.scale_bar.pack(side=Tkinter.LEFT,padx=15)

     F = Tkinter.Frame(Toolbar); F.pack(side=Tkinter.LEFT)
     self.enable_plotting = Tkinter.IntVar(); self.enable_plotting.set(1)
     if not hasattr(self,'make_plot'): self.enable_plotting.set(0)
     Tkinter.Checkbutton(F,text='Enable plotting',variable=self.enable_plotting).pack(side=Tkinter.TOP,anchor='w')
     
     self.enable_diagnostic = Tkinter.IntVar(); self.enable_diagnostic.set(1)
     Tkinter.Checkbutton(F,text='Enable diagnostic',variable=self.enable_diagnostic).pack(side=Tkinter.TOP,anchor='w')
     
     self.figure=Figure()
     self.figure.set_figwidth(float(self.image_width)/self.figure.get_dpi() )
     self.figure.set_figheight(float(self.image_height)/self.figure.get_dpi() )
     self.figure.text(0.5,0.5,self.pyOM_vstr,ha='center',va='center',fontsize='xx-large',color='darkblue')
     self._plotit()
     self.run_fortran('check_pyom_module') # check for consistent model variables
     self.itt  = 0; self.itt_str.set('time step %i' % self.itt)
     self.time = self.itt*M.dt
     self.diagnose()
     self.diag_plot()
     self.shell_on=False
     return


   def printplot(self):
     tkMessageBox.showinfo("Info", "not yet implemented")
     return

   def scale_both(self):
      if self.snapint == self.plotint:
         self.plotint = Tkinter.IntVar(); self.plotint.set( self.snapint.get() )
         self.scale_bar.config(variable=self.plotint,state=Tkinter.NORMAL)
      else:
         self.scale_bar.config(variable=self.snapint,state=Tkinter.DISABLED)
         self.plotint = self.snapint
      return

   def print_text(self,s):
     self.time_step_list.config(state=Tkinter.NORMAL)
     self.time_step_list.insert(Tkinter.END,s)
     self.time_step_list.yview(Tkinter.MOVETO,1)
     self.time_step_list.config(state=Tkinter.DISABLED)
     return
  
   def quit(self):
     if tkMessageBox.askokcancel(title='Quit',message='Exit pyOM?'): Tkinter.Frame.quit(self)
     return
   
   def reset(self):
     if tkMessageBox.askokcancel(title='Reset',message='Reset integration?'):
        self.print_text('\nresetting model integration')
        M=self.fortran.pyom_module         # fortran module with model variables
        M.u[:]=0; M.v[:]=0; M.w[:]=0
        M.b[:]=0; M.eta[:]=0; M.p_full[:]=0; M.psi[:]=0; M.p_surf[:]=0; M.p_hydro[:]=0
        self.initial_conditions()
        self.itt  = 0; self.itt_str.set('time step %i' % self.itt)
        self.time = self.itt*M.dt
        self.halt=1
        self.diagnose()
        self.diag_plot()
     return
      
   def read(self):
      file=tkFileDialog.askopenfilename()
      if file:
          self.halt=1
          self.print_text('\nreading restart file %s' %file)
          self.read_restart(file)
      return
           
   def write(self):
      self.halt=1
      self.print_text('\nstopping model integration')
      file=tkFileDialog.asksaveasfilename(initialfile='restart.dta')
      if file:
         self.print_text('\nwriting restart to file %s' %file)
         self.write_restart(file)
      return
    
   def stop(self):
      if self.halt == 0:
        self.halt=1
        self.print_text('\nstopping model integration')
      return

   def run(self):
      if self.halt == 1:
        self.halt=0
        self.print_text('\ncontinue model integration')
        self.one_time_step()
      return

   def one_time_step(self):
     """
     enter a simple model time stepping loop
     """
     if not self.halt:
       self.itt = self.itt + 1
       self.itt_str.set('time step %i' % self.itt)
       M=self.fortran.pyom_module         # fortran module with model variables
       self.time = self.itt*M.dt
       self.time_step()
       if numpy.mod(self.itt,self.snapint.get()) == 0 : self.diagnose()
       if numpy.mod(self.itt,self.plotint.get()) == 0 : self.diag_plot()
       self.time_goes_by()
       self.after(1,self.one_time_step)
     return

   def diagnose(self):
     """ diagnose the model variables, might be extended by further diagnostic
     """
     if self.enable_diagnostic.get()==1:
       M=self.fortran.pyom_module         # fortran module with model variables
       s= '\ndiagnosing at %f s, itt=%i, pressure solver itt = (%i,%i)'%(self.time,self.itt,M.sor2d_itts,M.sor3d_itts)
       self.print_text(s)
       # do something else here
     return
  
   def _plotit(self):
     from matplotlib.backends.backend_agg import FigureCanvasAgg
     import matplotlib.backends.tkagg  as tkagg
     canvas = FigureCanvasAgg(self.figure)
     canvas.draw()
     if hasattr(self,'main_area_id'): self.main_area.delete(self.main_area_id)
     self.main_area_id=self.main_area.create_image(self.image_width/2,self.image_height/2,image=self._tkphoto)
     tkagg.blit(self._tkphoto, canvas.renderer._renderer, colormode=2)
     return
     
   def _plotit_old(self):
     from matplotlib.backends.backend_agg import FigureCanvasAgg
     import Image,ImageTk
     canvas = FigureCanvasAgg(self.figure)
     canvas.draw()
     self.Image=ImageTk.PhotoImage( Image.fromstring("RGB",canvas.get_width_height(), canvas.tostring_rgb()) )
     self.main_area.create_image(self.image_width/2,self.image_height/2,image=self.Image)
      
   def diag_plot(self):
     """ diagnose the model variables, might be extended by further diagnostic
     """
     M=self.fortran.pyom_module         # fortran module with model variables
     if hasattr(self,'make_plot') and self.enable_plotting.get()==1 :
        self.print_text('\nplotting at %f s'% self.time )
        self.make_plot()
        self._plotit()
     if not hasattr(self,'make_plot'):
        self.print_text('\ncannot plot: missing method make_plot')
     return
   
   def print_parameter(self):
      pass

   def config(self):
     if hasattr(self,'config_setup'):
        self.config_window=Tkinter.Toplevel()
        self.config_window.protocol('WM_DELETE_WINDOW', self.quit)
        self.config_frame = Tkinter.Frame(self.config_window)
        self.config_frame.pack()
        #Tkinter.Label(self.config_frame,text='Config').pack(side=Tkinter.TOP)
        self.config_setup() # add widgets to frame
     else:
        tkMessageBox.showinfo("Info", "method config_setup not found")
     return
   
   def shell(self):
      if self.shell_on:
        self.shell_window.destroy()
        self.shell_on=False
      else:   
        self.shell_window=Tkinter.Toplevel()
        self.shell_window.protocol('WM_DELETE_WINDOW', self.quit)

        fo=("Times",12)
        self.display=ScrolledText.ScrolledText(self.shell_window,width=80,font=fo)
        self.display.pack()
      
        self.command = Tkinter.StringVar()
        self.command.set('')
        if not hasattr(self,'commandlist'):
           self.commandlist = ['']
           self.command_pointer = 0

        self.entry = Tkinter.Entry(self.display,textvariable=self.command,width=77,font=fo,bd=0)
        self.entry.config(highlightthickness=0)
        self.display.config(state=Tkinter.NORMAL)
        self.display.insert(Tkinter.END,'\n Welcome to %s \n\n'% self.pyOM_vstr)
        self.display.insert(Tkinter.END,'>>')
        self.display.see(Tkinter.END)
        self.display.window_create(Tkinter.END,window=self.entry)
        self.display.config(state=Tkinter.DISABLED)
      
        self.entry.bind(sequence="<Return>", func=self.command_process)
        self.entry.bind(sequence="<Up>",     func=self.command_up)
        self.entry.bind(sequence="<Down>",   func=self.command_down)
        self.entry.focus_set()
        self.entry.icursor(Tkinter.END)
        self.shell_on=True
        if not hasattr(self,'command_locals'): self.set_command_locals()
      return

   def set_command_locals(self):
     M=self.fortran.pyom_module         # fortran module with model variables
     import matplotlib.pyplot as plt
     self.command_locals=locals()
     return 

   def command_up(self,args):
      self.command_pointer = max(0,self.command_pointer-1)
      self.command.set(self.commandlist[ self.command_pointer ] )
      self.entry.icursor(Tkinter.END)
 
   def command_down(self,args):
      self.command_pointer = min(len(self.commandlist),self.command_pointer+1 )
      if self.command_pointer == len(self.commandlist):
       self.command.set('')
      else:
       self.command.set( self.commandlist[ self.command_pointer ] )
      self.entry.icursor(Tkinter.END)
      
   def command_process(self, args):
      import sys,StringIO
      self.sendToDisplay('>>'+self.command.get()+"\n")
      if self.command.get() != '':
         self.commandlist.append(self.command.get() )
         self.command_pointer = len(self.commandlist )
      b1=sys.stderr; b2=sys.stdout
      sys.stderr = StringIO.StringIO()
      sys.stdout = StringIO.StringIO()
      c= code.InteractiveInterpreter(self.command_locals)
      c.runsource(self.command.get() )
      self.sendToDisplay(sys.stderr.getvalue())
      self.sendToDisplay(sys.stdout.getvalue())
      sys.stderr=b1; sys.stdout=b2
      self.command.set('')
      return
   
   def sendToDisplay(self, string):
      self.display.config(state=Tkinter.NORMAL)
      self.display.insert(Tkinter.END+'-4 chars', string)
      self.display.see(Tkinter.END)
      self.display.config(state=Tkinter.DISABLED)
      return
   

