#!/usr/bin/python
"""
 Python interface to the IFEFFIT XAFS Analysis library

 This module provides a simple set of functions to interface with
 the IFEFFIT XAFS Analysis library. The functions provided are:

   - ifeffit(cmd):         execute an Ifeffit command
   - put_scalar(name,val): set a python double as an Ifeffit scalar
   - get_scalar(name):     return an Ifeffit scalar as a python double
   - put_string(name,val): set a python string as an Ifeffit string
   - get_string(name):     return an Ifeffit string as a python string
   - get_string(name):     return an Ifeffit string as a python string
   - put_array(name,list): set a python list as an Ifeffit array
   - get_echo():           get lines from Ifeffit's echo cache
   - shell():              execute a simple shell interface for Ifeffit


 M Newville Univ of Chicago
"""
##
##  Ifeffit python module: XAFS Analysis Library
##
##  Copyright (c) 1997--2000  Matthew Newville, The University of Chicago
##  Copyright (c) 1992--1996  Matthew Newville, University of Washington
##
##  Permission to use and redistribute the source code or binary forms of
##  this software and its documentation, with or without modification is
##  hereby granted provided that the above notice of copyright, these
##  terms of use, and the disclaimer of warranty below appear in the
##  source code and documentation, and that none of the names of The
##  University of Chicago, The University of Washington, or the authors
##  appear in advertising or endorsement of works derived from this
##  software without specific prior written permission from all parties.
##
##  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 THIS SOFTWARE.
##


import _ifeffit, types, cmd, os, sys
from string import strip,find, split

__version__  = '1.2'

##{ generated by swig
def _swig_setattr(self,class_type,name,value):
    if (name == "this"):
        if isinstance(value, class_type):
            self.__dict__[name] = value.this
            if hasattr(value,"thisown"): self.__dict__["thisown"] = value.thisown
            del value.thisown
            return
    method = class_type.__swig_setmethods__.get(name,None)
    if method: return method(self,value)
    self.__dict__[name] = value

def _swig_getattr(self,class_type,name):
    method = class_type.__swig_getMethods__.get(name,None)
    if method: return method(self)
    raise AttributeError,name

try:
    _object = types.ObjectType
    _newclass = 1
except AttributeError:
    class _object : pass
    _newclass = 0

##}

if (os.name == 'nt'):
    oskeys  = os.environ.keys()
    defkeys = {'PGPLOT_DEV':'/GW',
               'IFEFFIT_DIR':'C:\Program Files\Ifeffit'}
    for i in defkeys.keys():
        if (i in oskeys):
            u = os.environ[i]
            v = defkeys[i]
        else:
            os.environ[i] =  defkeys[i]
    os.environ['PGPLOT_DIR']  = os.environ['IFEFFIT_DIR']
    os.environ['PATH']  = os.environ['PATH']  + ';.;' + os.environ['IFEFFIT_DIR']

try:
    import _ifeffit
    _ifeffit.iff_exec(' ')
except:
    print "Cannot load Ifeffit Module -- check path and installation"
    sys.exit(1)



class Ifeffit:
    """
 Ifeffit -  Python interface to the IFEFFIT XAFS Analysis library

 This module provides a simple set of functions to interface with
 the IFEFFIT XAFS Analysis library. The functions provided are:

   - ifeffit(cmd):         execute an Ifeffit command
   - put_scalar(name,val): set a python double as an Ifeffit scalar
   - get_scalar(name):     return an Ifeffit scalar as a python double
   - put_string(name,val): set a python string as an Ifeffit string
   - get_string(name):     return an Ifeffit string as a python string
   - get_string(name):     return an Ifeffit string as a python string
   - put_array(name,list): set a python list as an Ifeffit array
   - get_echo():           get lines from Ifeffit's echo cache
   - shell():              execute a simple shell interface for Ifeffit


 M Newville Univ of Chicago
"""
    MAX_ARRAY_PTS = 16384
    def __init__(self,screen_echo = 1,use_numeric=1):
        # this initializes the Ifeffit session
        try:
            _ifeffit.iff_exec(' ')
        except:
            print " failed to load ifeffit library"
            sys.exit(1)

        startup_string = "&screen_echo=%i" % (screen_echo)
        self.ifeffit(startup_string)
        self.MAX_ARRAY_PTS = int(self.get_scalar("&maxpts"))
        self.num_array = None
        if (use_numeric == 1):
            try:
                import Numeric
                self.num_array = Numeric.array
            except:
                self.num_array = None

    def show_numeric(self):
        return (self.num_array != None)

    def ifeffit(self,cmd):
        "execute an Ifeffit command or list of ifeffit commands"
        coms = split(cmd,'\012')
        ret  = 0
        for c in coms:
            if (c != ''):  ret  = _ifeffit.iff_exec(c)
        return ret

    def __call__(self,cmd):  self.ifeffit(cmd)
        
    def put_scalar(self,name,val):
        "set a python double as an Ifeffit scalar"
        ret = self.ifeffit("%s = %s" % (name, val))
        return ret

    def get_scalar(self,name):
        "return an Ifeffit scalar as a python double"
        ptr = _ifeffit.new_Pdbl()
        u   = _ifeffit.iff_get_scalar(name, ptr)
        val = _ifeffit.Pdbl_value(ptr)
        _ifeffit.delete_Pdbl(ptr)
        return val

    def put_string(self,name,val):
        "set a python string as an Ifeffit string"
        ret = self.ifeffit("$%s = %s" % (name, val))
        return ret

    def get_string(self,name):
        "return an Ifeffit string as a python string"
        str = " "*256
        _ifeffit.iff_get_string(name, str)
        ret = strip(str)
        return ret

    def get_echo(self):
        "get lines from Ifeffit's echo cache."
        str = " "*512
        ret = _ifeffit.iff_get_echo(str)
        return str[0:ret]

    def get_echo_buffer(self):
        "return the full echo buffer"
        s = []
        for i in range(int(self.get_scalar('&echo_lines'))):
            x = strip(self.get_echo())
            if (x != ''): s.append(x)
        return s

    def clear_echo_buffer(self):
        "make sure IFeffit's echo buffer is clean"
        s = self.get_echo_buffer()
        return None

    def clear_arrays(self):
        m = self.list_arrays()
        for i in m:
            self.ifeffit("erase %s" % i)
            
    def list_arrays(self):
        self.get_echo_buffer()
        echo_orig =  int(self.get_scalar("&screen_echo"))
        self.ifeffit("set &screen_echo=0")
        self.ifeffit("show @arrays")
        s = self.get_echo_buffer()
        x = []
        for i in s:  x.append(i.split('=')[0].strip())
        self.ifeffit("set &screen_echo=%i"%echo_orig)            
        return x
    
        
    def put_array(self,name,data):
        "set a python list as an Ifeffit array"
        if (self.num_array != None):   data = list(data)
        npts  = min(self.MAX_ARRAY_PTS, len(data))
        ptr   = _ifeffit.new_Parr(npts)
        for i in range(npts):
            _ifeffit.Parr_setitem(ptr,i,data[i])

        p_n   = _ifeffit.new_Pint()
        _ifeffit.Pint_assign(p_n,npts)

        ret = _ifeffit.iff_put_array(name, p_n, ptr)
        _ifeffit.delete_Pint(p_n)
        _ifeffit.delete_Parr(ptr)
        return ret

    def get_array(self,name):
        "return an Ifeffit array as a python list"
        ptr  = _ifeffit.new_Parr(self.MAX_ARRAY_PTS)
        npts = _ifeffit.iff_get_array(name,ptr)
        arr  = []
        for i in range(npts):
            arr.append(_ifeffit.Parr_getitem(ptr,i))
        _ifeffit.delete_Parr(ptr)
        if (self.num_array != None):  arr = self.num_array(arr)
        return arr


    def read_file(self,file=None):
        "return data structure for arrays read from a file"
        if (file == None): return None
        self.ifeffit("read_data(%s, group=t)" % (file))
        tmp = self.get_string("column_label")
        col = string.split(tmp)
        arr = []
        for i in col:
            arr.append(iff.get_array("t.%s" % (i)))
        if (self.num_array != None):  arr = self.num_array(arr)
        return (arr,col)

    def get_group_arrays(self):
        "return dictionary of ifeffit groups/aarrays"
        self.clear_echo_buffer()
        dict = {}
        # print ' IFF get_group_arrays!'
        self.ifeffit("show @arrays");
        buff= self.get_echo_buffer()
        for s in buff:
            if ((s != "") and (find(s, '=') != -1)):
                sx   = split(s,'=')
                npts = 0
                if (len(sx) < 0):  sx = [' ',' ']
                npts = strip(split(sx[1],'pts')[0])
                name = sx[0]
                try:
                    (pre,suff)  = split(name,'.')
                    if (dict.has_key(pre) == 0):
                        dict[pre] = {'names':[], 'npts':[]}
                    dict[pre]['names'].append(strip(suff))
                    dict[pre]['npts'].append(npts)
                except ValueError:
                    pass
        return dict
    #
    def reset_plot_opts(self):
        " set generic plotting options into structure"
        p   = {'bg': 'white', 'fg': 'black',
               'gridcolor': '#CCBEE0',
               'showgrid' : 1,   'charsize' : 1.5,
               'charfont' : 1,   'linewidth': 2,
               'linestyle': 1,
               'title': "",
               'xlabel': "",  'ylabel': "",
               'xmin': "",    'xmax': "",
               'ymin': "",    'ymax': "",
               'traces': [['blue',      'solid'],
                          ['blue',      'solid'],
                          ['red',       'solid'],
                          ['black',     'solid'],
                          ['darkgreen', 'solid'],
                          ['magenta',   'solid'],
                          ['maroon',    'solid'],
                          ['yellow',    'solid'],
                          ['orange',    'solid'],
                          ['purple',    'solid']]
               }
        return p


    def get_plot_opts(self):
        " get plotting options from current ifeffit session into structure"
        p   = self.reset_plot_opts()
        # print "getting plot opts"
        self.clear_echo_buffer()
        self.ifeffit('color("show")')
        nt = len(p['traces'])
        x  = self.get_echo()
        for i in range(self.get_scalar('&echo_lines')):
            x = strip(self.get_echo())
            x = split(x, '=')
            k = strip(x[0])
            v = strip(x[1])
            if (k == 'bg'):
                p['bg'] = v
            elif (k == 'fg'):
                p['fg'] = v
            elif (k == 'grid'):
                p['gridcolor'] = v
            else:
                try:
                    j = int(k)
                    if (j < nt):
                        p['traces'][j][0] = v
                    else:
                        p['traces'].append([ v, 'solid'])
                except ValueError:
                    pass
    ##
        self.ifeffit('linestyle("show")')
        n_traces = len(p['traces'])
        x = self.get_echo()
        for i in range(self.get_scalar('&echo_lines')):
            x = strip(self.get_echo())
            x = split(x, '=')
            k = strip(x[0])
            v = strip(x[1])
            try:
                j = int(k)
                if (j < n_traces):
                    p['traces'][j][1] = v
                else:
                    p['traces'].append([ 'black', v])
            except ValueError:
                pass
    ##
        return p

    def set_plot_opts(self,p):
        " set plotting to ifeffit session from plot structure"
        self.ifeffit ("color(fg='%s', bg='%s', gridcolor='%s')"%
                      (p['fg'], p['bg'], p['gridcolor']) )

        for i in range(1,len(p['traces'])):
            self.ifeffit ("color(%i='%s')"     % ( i, p['traces'][i][0]) )
            self.ifeffit ("linestyle(%i='%s')" % ( i, p['traces'][i][1]) )


class iff_shell(cmd.Cmd,Ifeffit):
    __doc__= """
  IFeffit Shell in Python

  This simple Ifeffit Shell provides enough functionality for most
  use.  Command history and editing are supported (with the arrow
  keys and other GNU Readline mechanisms). Shell commands will be
  executed if preceded by a '!' (and not a whimper!).

  For a  more functional command shell see the C version

  M Newville  Univ of Chicago
"""
    ps1    = "Ifeffit>"
    ps2    = "  ...  >"
    intro  = " Ifeffit Shell in Python "
    misc_header = "Help is also available on: "
    max_save_lines = 500
    def __init__(self):
        import os
        try:
            import readline
            self.rdline = readline
            import rlcompleter
            self.rdline.parse_and_bind("tab: complete")
        except ImportError:
            self.rdline = None

            
        Ifeffit.__init__(self)
        self.prompt    = self.ps1
        return None

    def __del__(self):
        if (self.rdline):
            self.rdline.write_history_file('iff_shell.iff')

    def emptyline(self):
        pass

    def help_help(self):
        print self.__doc__
    def help_Ifeffit(self):
        print Ifeffit.__doc__
    def help_shell(self):
        print """
        shell commands can be executed by preceding either by 'shell' or '!':
   Ifeffit>! ls
   Ifeffit>shell lpr ifeffit.ps"""
    def help_EOF(self):
        print " EOF = quit = exit = Control^D ==exit the program"
    def do_shell(self, arg):
        import os
        os.system(arg)
    def do_EOF(self,arg):
        return 1
    def default(self,str1):
        str = strip(str1)
        if ((str == 'quit') or (str == 'exit')):
            return 1
        elif (str[0:1] == '!'):
            self.do_shell(str)
        else:
            j = self.ifeffit(str)
            self.prompt = self.ps1
            if (j == -1): self.prompt = self.ps2


#
# if this was invoked from a command line, execute the shell!
if (__name__ == '__main__'):
    i =  iff_shell()
    i.cmdloop()

