#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright (c) 2012-2013, tamanegi (tamanegi@users.sourceforge.jp)
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.

import wx
from wx import xrc

import sys
import platform
import re
import os
import StringIO

filedirname = os.path.dirname( os.path.abspath( __file__ ) )
commonpath = os.path.join( filedirname, "../common" )
if not commonpath in sys.path:
    sys.path.append( commonpath )
from vector3d import *
from pdb import *
from system import *
from sysmgr import *

XRC_FILE = "wxbuilder.xrc"

class WMBuilderWX( wx.App ):
  ### functions involving constuctor and its sub-functions
  def OnInit( self ):
    self.res = xrc.XmlResource( XRC_FILE )
    self._initializeBuilderVars()
    self._initializeMainFrame()
    self._initializeSubFrames()
    self._defineEvents()
    self._defineSelectionEvents()
    self._defineXmlEvents()
    return True

  def _initializeSelectionStates( self ):
    # whether all atoms are selected
    self.__state_selectAll = True
    # whether no atoms are selected
    self.__state_selectNone = False
    # if none of above flags are on, some of residues/atoms might be selected

  def _initializeBuilderVars( self ):
    self.smgr = SystemMgr()

  def _initializeMainFrame( self ):
    self.frameMain = self.res.LoadFrame( None, "MainFrame" )
    self.frameMain.SetTitle( "wxWMBuilder" )
    self._initializeMenu()
    self.frameMain.Show()

  def _initializeMenu( self ):
    self.MenuMain = self.frameMain.GetMenuBar()
    self.MenuWindow = self.MenuMain.GetMenu( self.MenuMain.FindMenu( "Window" ) )

  def _initializeSubFrames( self ):
    # selection
    self.buttons = []
    self.layouts = []
    self.frameSubSelection = self.res.LoadFrame( None, "SelectionFrame" )
    self.frameSubSelection.SetTitle( "Selection Manager" )
    self.miscSetSelectionWindowButtonsState()
    self._initializeSelectionStates()
    # xml
    self.xmllayout = None
    self.frameSubXml = self.res.LoadFrame( None, "XmlFrame" )

  ### define events
  # main window
  def _defineEvents( self ):
    # close events
    self.frameMain.Bind( wx.EVT_CLOSE, self.evMenuFileQuit )
    # button events
    self.Bind( wx.EVT_BUTTON, self.evFileRefButton, id=xrc.XRCID("FileRefButton") )
    self.Bind( wx.EVT_BUTTON, self.evLoadButton, id=xrc.XRCID("LoadButton") )
    # menu events
    self.Bind( wx.EVT_MENU, self.evMenuFileOpen, id=xrc.XRCID("MenuFileOpen") )
    self.Bind( wx.EVT_MENU, self.evMenuFileQuit, id=xrc.XRCID("MenuFileQuit") )
    self.Bind( wx.EVT_MENU, self.evMenuWindowSelection, id=xrc.XRCID("MenuWindowSelection") )
    self.Bind( wx.EVT_MENU, self.evMenuWindowXml, id=xrc.XRCID("MenuWindowXml") )

  # selection window
  def _defineSelectionEvents( self ):
    # close events
    self.frameSubSelection.Bind( wx.EVT_CLOSE, self.evCloseSelection )
    # button events
    self.Bind( wx.EVT_BUTTON, self.evSelectAllButton, id=xrc.XRCID("SelectAllButton") )
    self.Bind( wx.EVT_BUTTON, self.evClearSelectButton, id=xrc.XRCID("ClearSelectButton") )
    self.Bind( wx.EVT_BUTTON, self.evAddQueryButton, id=xrc.XRCID("AddQueryButton") )
    self.Bind( wx.EVT_BUTTON, self.evClearQueryButton, id=xrc.XRCID("ClearQueryButton") )
    self.Bind( wx.EVT_BUTTON, self.evApplySelectButton, id=xrc.XRCID("ApplySelectionButton") )
    self.Bind( wx.EVT_BUTTON, self.evClearSelectedAtomsButton, id=xrc.XRCID("ClearSelectedAtomsButton") )

  # Xml Window
  def _defineXmlEvents( self ):
    # close events
    self.frameSubXml.Bind( wx.EVT_CLOSE, self.evCloseXml )
    # button events
    self.Bind( wx.EVT_BUTTON, self.evXmlGenAtomsButton, id=xrc.XRCID("GenXmlAtomsButton") )
    self.Bind( wx.EVT_BUTTON, self.evXmlGenBondsButton, id=xrc.XRCID("GenXmlBondsButton") )
    self.Bind( wx.EVT_BUTTON, self.evXmlGenRibbonsButton, id=xrc.XRCID("GenXmlRibbonsButton") )
    self.Bind( wx.EVT_BUTTON, self.evXmlGenCoilsButton, id=xrc.XRCID("GenXmlCoilsButton") )
    self.Bind( wx.EVT_BUTTON, self.evUpdateXmlButton, id=xrc.XRCID("UpdateXmlButton") )
    self.Bind( wx.EVT_BUTTON, self.evClearXmlButton, id=xrc.XRCID("ClearXmlButton") )
    self.Bind( wx.EVT_BUTTON, self.evSaveXmlToFileButton, id=xrc.XRCID("SaveXmlToFileButton") )

  ### event handlers; main window
  # quit program
  def evQuit( self, event = None ):
    self.Destroy()
    sys.exit()

  # close events
  def evCloseSelection( self, event ):
    self.frameSubSelection.Hide()
    mi = self.MenuWindow.FindItemById( xrc.XRCID("MenuWindowSelection") )
    mi.Check( False )

  def evCloseXml( self, event ):
    self.frameSubXml.Hide()
    mi = self.MenuWindow.FindItemById( xrc.XRCID("MenuWindowXml") )
    mi.Check( False )

  # show file dialog to choose file
  def evFileRefButton( self, event = None ):
    dialog = wx.FileDialog( self.frameMain, "Choose a PDB file", ".", "", "*.pdb", wx.FD_OPEN )
    try:
      if dialog.ShowModal() == wx.ID_OK:
        tctl = xrc.XRCCTRL( self.frameMain, "textCtrlFilename" )
        tctl.SetValue( dialog.GetFilename() )
    finally:
      dialog.Destroy()

  def evLoadButton( self, event = None ):
    tctl = xrc.XRCCTRL( self.frameMain, "textCtrlFilename" )
    try:
      fh = open( tctl.GetValue(), "r" )
    except IOError, ( errno, msg ):
      if len( tctl.GetValue() ) == 0:
        self.miscShowMessageBox( "Please input filename" )
      else:
        self.miscShowMessageBox( "Unable to open file " + tctl.GetValue() )
      return
    self.smgr.clear()
    self.smgr.readPDB( fh )
    tctl = xrc.XRCCTRL( self.frameMain, "StatNumAtoms" )
    tctl.SetLabel( str(self.smgr.totatom()) )
    # create residue buttons
    self._loadButtonsSubSelection()
    # reset selection states
    self._initializeSelectionStates()
    # reset button state
    self.miscSetSelectionWindowButtonsState()
    self.frameSubSelection.Refresh()
    self.frameSubSelection.Update()
    fh.close()

  def evMenuFileOpen( self, event ):
    self.evFileRefButton()
    self.evLoadButton()

  def evMenuFileQuit( self, event ):
    dialog = wx.MessageDialog( self.frameMain, "Really quit?", style=wx.YES_NO|wx.NO_DEFAULT )
    try:
      if dialog.ShowModal() == wx.ID_YES:
        self.evQuit()
    finally:
      dialog.Destroy()

  def evMenuWindowSelection( self, event ):
    # MenuItem cannot be retrieved by XRCCTRL, since MenuItem does not derive
    # from Window class.
    mi = self.MenuWindow.FindItemById( xrc.XRCID("MenuWindowSelection") )
    if mi.IsChecked():
      self.frameSubSelection.Show()
    else:
      self.frameSubSelection.Hide()

  def evMenuWindowXml( self, event ):
    # see evMenuWindowSelection above
    mi = self.MenuWindow.FindItemById( xrc.XRCID("MenuWindowXml") )
    if mi.IsChecked():
      self.frameSubXml.Show()
    else:
      self.frameSubXml.Hide()

  ### event handlers; selection window
  def evResidueButton( self, event = None ):
    # just reset flags. There is no specific function for each button.
    self.__state_selectAll = False
    self.__state_selectNone = False

  def evSelectAllButton( self, event = None ):
    self.enableToggleButtons()
    self._initializeSelectionStates()

  # clear current selection; already registered querys do not change
  def evClearSelectButton( self, event = None ):
    self.disableToggleButtons()
    self.__state_selectAll = False
    self.__state_selectNone = True

  def evApplySelectButton( self, event ):
    tctl = xrc.XRCCTRL( self.frameSubSelection, "textCtrlQueryText" )
    # if string field is not empty, do not apply selection
    if tctl.GetValue():
      self.miscShowMessageBox( "Please click \"Add Query\" or remove text in the TextField above." )
      return
    if not self.__state_selectAll:
      if self._isToggleButtonEnabled():
        self.miscShowMessageBox( "Please click \"Add Query\" or disable toggle buttons." )
        return
    self.smgr.select()
    self.evClearSelectButton()
    # show number of selected atoms in the main window
    tctl = xrc.XRCCTRL( self.frameSubSelection, "StatNumSelectedAtoms" )
    tctl.SetLabel( str(self.smgr.natom()) )

  def evClearSelectedAtomsButton( self, event ):
    # clear currnt selected atominformation
    for each_sys in self.smgr.systems:
      each_sys.data = []
      each_sys.rawdata = []
    # reset label
    tctl = xrc.XRCCTRL( self.frameSubSelection, "StatNumSelectedAtoms" )
    tctl.SetLabel( "" )

  def evAddQueryButton( self, event ):
    # check residue buttons
    num = 0
    if not self.__state_selectAll:
      for each_sys in self.smgr.systems:
        for each_mol in each_sys.PDB.molecules:
          for button in self.buttons[num]:
            if not button.GetValue():
              continue
            res = each_mol.residues.get( button.GetLabel() )
            for atom in res.atoms:
              each_sys.PDB.rawdata.append( atom )
          num += 1
      # check text area
      tctl = xrc.XRCCTRL( self.frameSubSelection, "textCtrlQueryText" )
      tx = tctl.GetValue() # query text
      if len( tx.strip() ) > 0:
        # check radio button state
        rb = xrc.XRCCTRL( self.frameSubSelection, "radioBoxQuery" )
        v = rb.GetSelection()
        if v == 0:
          self.smgr.addMolQuery( tx )
        elif v == 1:
          self.smgr.addResQuery( tx )
        elif v == 2:
          self.smgr.addAtomQuery( tx )
        tctl.SetValue("")
    # reset query
    self.evClearSelectButton()

  # remove all querys
  def evClearQueryButton( self, event ):
    self.smgr.initQuery()
    # clear rawdata
    for each_sys in self.smgr.systems:
      each_sys.PDB.rawdata = []
    self._initializeSelectionStates()
    self.disableToggleButtons()

  ### disable buttons
  def disableButtonsSelection( self ):
    self.toggleApplySelectionButton( False )
    self.toggleSelectAllButton( False )
    self.toggleClearSelectButton( False )

  def enableButtonsSelection( self ):
    self.toggleApplySelectionButton( True )
    self.toggleSelectAllButton( True )
    self.toggleClearSelectButton( True )

  def toggleApplySelectionButton( self, val ):
    ab = xrc.XRCCTRL( self.frameSubSelection, "ApplySelectionButton" )
    ab.Enable( val )

  def toggleSelectAllButton( self, val ):
    sa = xrc.XRCCTRL( self.frameSubSelection, "SelectAllButton" )
    sa.Enable( val )

  def toggleClearSelectButton( self, val ):
    cs = xrc.XRCCTRL( self.frameSubSelection, "ClearSelectButton" )
    cs.Enable( val )

  def toggleAddQueryButton( self, val ):
    aq = xrc.XRCCTRL( self.frameSubSelection, "AddQueryButton" )
    aq.Enable( val )

  def toggleClearQueryButton( self, val ):
    cq = xrc.XRCCTRL( self.frameSubselection, "ClearQueryButton" )
    cq.enable( val )

  ### event handlers and related functions for XML window
  def evXmlGenAtomsButton( self, event ):
    nelem = self.smgr.genAtoms()
    self.miscShowNumberGeneratedXmlElements( nelem )

  def evXmlGenBondsButton( self, event ):
    nelem = self.smgr.genBonds()
    self.miscShowNumberGeneratedXmlElements( nelem )

  def evXmlGenRibbonsButton( self, event ):
    nelem = self.smgr.genChains()
    self.miscShowNumberGeneratedXmlElements( nelem )

  def evXmlGenCoilsButton( self, event ):
    nelem = self.smgr.genChains( True )
    self.miscShowNumberGeneratedXmlElements( nelem )

  def evUpdateXmlButton( self, event ):
    if self.xmllayout != None:
      self.xmllayout.DeleteWindows()
    # get xml data
    self.smgr.addwmxml = True
    tmpout = StringIO.StringIO()
    self.smgr.purgeData( tmpout )
    xmlstring = tmpout.getvalue()
    tmpout.close()
    # upload string to the window
    sw = xrc.XRCCTRL( self.frameSubXml, "XMLDataWindow" )
    sw.SetScrollbars( 5, 5, 1, 1 )
    self.xmllayout = wx.BoxSizer( wx.VERTICAL )
    strings = xmlstring.split('\n')
    for mystr in strings:
      self.xmllayout.Add( wx.StaticText( sw, wx.ID_ANY, label=mystr ) )
    sw.SetSizer( self.xmllayout )
    sw.FitInside()
    sw.Refresh()
    sw.Update()

  def evClearXmlButton( self, event ):
    for each_sys in self.smgr.systems:
      each_sys.scenedata = ""

  def evSaveXmlToFileButton( self, event ):
    dialog = wx.FileDialog( self.frameSubXml, "Choose output file name", ".", "", "*.xml", wx.FD_SAVE )
    try:
      if dialog.ShowModal() == wx.ID_OK:
        filename = dialog.GetFilename()
        fh = open( filename, "w" )
    except IOError, ( errno, msg ):
      # maybe failed to open file
      if len( filename ) == 0:
        self.miscShowMessageBox( "Please input output filename." )
      else:
        self.miscShowMessageBox( "Unable to open file " + filename + " for writing." )
      return
    finally:
      dialog.Destroy()
    # write to file
    self.smgr.purgeData( fh )
    fh.close()

  ### misc functions
  def miscSetSelectionWindowButtonsState( self ):
    # if no structure data is loaded, disbale all buttons
    if self.smgr.totatom() == 0:
      self.disableButtonsSelection()
    else:
      self.enableButtonsSelection()

  # show message box (may be removed later)
  def miscShowMessageBox( self, msg ):
    wx.MessageBox( msg, 'Info' )
    
  def miscShowNumberGeneratedXmlElements( self, nelem ):
    tctl = xrc.XRCCTRL( self.frameSubXml, "StatNumGeneratedXmlElements" )
    tctl.SetLabel( str(nelem) + " XML elements generated" )

  ### dynamically created toggle buttons
  def enableToggleButtons( self ):
    for buttons in self.buttons:
      for button in buttons:
        button.SetValue( True )
    self.miscSetSelectionWindowButtonsState()

  def disableToggleButtons( self ):
    for buttons in self.buttons:
      for button in buttons:
        button.SetValue( False )
    self.miscSetSelectionWindowButtonsState()

  def _isToggleButtonEnabled( self ):
    for buttons in self.buttons:
      for button in buttons:
        if button.GetValue():
          return( True )
    return( False )

  ### selection buttons
  def _loadButtonsSubSelection( self ):
    self._destroyOldLayoutsAndButtons()
    sw = xrc.XRCCTRL( self.frameSubSelection, "GUISelection" )
    sw.SetScrollbars( 5, 5, 1, 1 )
    self.layouts = []
    self.buttons = []
    mol_id = -1
    # number of residues for each molecules
    for each_sys in self.smgr.systems:
      for each_mol in each_sys.PDB.molecules:
        mol_id += 1
        # use a panel for each molecule
        layout = wx.BoxSizer( wx.HORIZONTAL )
        buttons = []
        res_id = -1
        for reskey in each_mol.reskeys:
          if len( buttons ) != 0:
            # will be replaced by a spacer of equal width
            layout.Add( wx.StaticLine( sw, size=(20,0) ) )
          res_id += 1
          button = wx.ToggleButton( sw, wx.ID_ANY, reskey, style=wx.BU_EXACTFIT )
          # common event; this events just handle global flags
          self.Bind( wx.EVT_TOGGLEBUTTON, self.evResidueButton, button )
          button.SetValue( True )
          buttons.append( button )
          layout.Add( buttons[-1] )
        self.buttons.append( buttons )
        self.layouts.append( layout )
        # do not use SetSizerAndFit; scrollbar won't be shown
        sw.SetSizer( self.layouts[-1] )
        sw.FitInside()
    sw.Refresh()
    sw.Update()

  # clear all buttons
  def _destroyOldLayoutsAndButtons( self ):
    for layout in self.layouts:
      layout.DeleteWindows()

if __name__ == "__main__":
  app = WMBuilderWX( False )
  app.MainLoop()
