#
# Licensing notes:  This is released under the PyMOL license.
#
# TODO:
# 
from pymol.wizard import Wizard
from pymol import cmd
import sys,os,struct

######
# Some fixes to ensure Python 2.2 compatibility
######

try:
    junk = enumerate
except:
    def enumerate(collection):
        '''Generates an indexed series:  (0,coll[0]), (1,coll[1]) ...

        Taken directly from PEP 279 http://www.python.org/peps/pep-0279.html
        '''     
        i = 0
        it = iter(collection)
        while 1:
            yield (i, it.next())
            i += 1
try:
    junk = True
except:
    True = 1 == 1
    False = not True

defaults = {
    "iso_pos" : None,
    "iso_neg" : None,
    "ramp_low" : -0.1,
    "ramp_middle" : 0.0,
    "ramp_high" : 0.1,
    "showHelpPrompt" : False,
    "vdw_modifier" : None,
    }

class Electrostatics(Wizard):

    def __init__(self):
        Wizard.__init__(self)

        for k,v in defaults.iteritems():
            self.set(k,v,False)

        self.helpPrompt = self.getHelpPrompt()

        # self.avail_maps is all available maps
        # self.map is the one we're using
        self.map = None
        self.avail_maps = []
        for a in cmd.get_names(): # automatically load the first map that we find
            if cmd.get_type(a)=='object:map':
                self.map = a
                break

        self.molecule = None
        self.avail_molecules = []
        for a in cmd.get_names(): # automatically pick the first molecule we find
            if cmd.get_type(a)=='object:molecule':
                self.molecule = a
                break
            


    def set(self,attribute,value,refresh=True):
        setattr(self,attribute,value)
        if refresh:
            cmd.refresh_wizard()

    def getHelpPrompt(self):
        """Get the help prompt
        """
        helpMessage="""
If you have a dx map already, you can load it with 'load_dx map.dx <name>'
"""
        return helpMessage.split('\n')[1:]
        
    def get_panel(self):
        # 0 is blank
        # 1 is text
        # 2 is button
        # 3 is popup

        #
        # We need to keep a lot of this here instead of in __init__ so that
        # it will be re-calculated each time we need it.  It's easiest to
        # have everything in one place, so all of self.menu goes here.
        #

        self.menu['map'] = [[1,a,'cmd.get_wizard().set("map","%s")'%a] for a in cmd.get_names('objects') if cmd.get_type(a)=='object:map']

        self.menu['molecule'] = [[1,a,'cmd.get_wizard().set("molecule","%s")'%a] for a in cmd.get_names('objects') if cmd.get_type(a)=='object:molecule']

        self.menu['isosurf'] = [[1,'%.2f'%v,'cmd.get_wizard().isosurf(%.2f)'%v ] for v in (0.0,-0.1,0.1,-1.0,1.0)] + [[1, 'None','cmd.get_wizard().isosurf(None)'],]

        self.menu['isosurf_pos'] = [[1,'%.2f'%v,'cmd.get_wizard().isosurf(%.2f,"iso_pos")'%v ] for v in (0.0,0.1,0.5,1.0,2.0,5.0,10.0,15.0)] + [[1, 'None','cmd.get_wizard().isosurf(None,"iso_pos")'],]

        self.menu['isosurf_neg'] = [[1,'%.2f'%v,'cmd.get_wizard().isosurf(%.2f,"iso_neg")'%v ] for v in (0.0,-0.1,-0.5,-1.0,-2.0,-5.0,-10.0,-15.0)] + [[1, 'None','cmd.get_wizard().isosurf(None,"iso_neg")'],]

        self.menu['surface'] = [[1,'%.2f'%v,'cmd.get_wizard().surface(%.2f)'%v ] for v in (0.0,1,2,3,4,5,10,15)] + [[1, 'None','cmd.get_wizard().surface(None)'],]

        self.menu['transparency'] = [[1, '%.1f'%v,'cmd.set("transparency",%f,"%s")'%(v,self.molecule)] for v in (0,.1,.2,.3,.4,.5,.6,.7,.8,.9)]

        ramp_range = [self.ramp_low + i for i in (-5,-3,-1,-0.5,-0.3,-0.1,-0.05,-0.03,-0.01,0,0.01,0.03,0.05,0.1,0.3,0.5,1,3,5)]
        self.menu['ramp_low'] = [[1,str(v),'cmd.get_wizard().ramp(low=%.2f)'%v] for v in ramp_range]

        ramp_range = [self.ramp_middle + i for i in (-5,-3,-1,-0.5,-0.3,-0.1,-0.05,-0.03,-0.01,0,0.01,0.03,0.05,0.1,0.3,0.5,1,3,5)]

        self.menu['ramp_middle'] = [[1,str(v),'cmd.get_wizard().ramp(middle=%.2f)'%v] for v in ramp_range]

        ramp_range = [self.ramp_high + i for i in (-5,-3,-1,-0.5,-0.3,-0.1,-0.05,-0.03,-0.01,0.0,0.01,0.03,0.05,0.1,0.3,0.5,1,3,5)]
        self.menu['ramp_high'] = [[1,str(v),'cmd.get_wizard().ramp(high=%.2f)'%v] for v in ramp_range]


        result = [[ 1, 'Electrostatics',''],]
        if self.showHelpPrompt:
            result += [[2, 'Hide help','cmd.get_wizard().hide_help()'],]
        else:
            result += [[2, 'Show help','cmd.get_wizard().show_help()'],]
        result += [
            [ 3, 'Molecule (%s)'%self.molecule,'molecule'],
            [ 3, 'Map (%s)'%self.map,'map'],
            [ 1, 'Isosurfaces',''],
            [ 3, '  Positive %s'%self.iso_pos,'isosurf_pos'],
            [ 3, '  Negative %s'%self.iso_neg,'isosurf_neg'],
            [ 1, 'Molecular Surface',''],
            [ 3, '  Vdw Buffer (%s)'%self.vdw_modifier,'surface'],
            [ 3, '  Transparency','transparency'],
            [ 1, 'Ramp',''],
            [ 3, '  Low %.3f'%self.ramp_low,'ramp_low'],
            [ 3, '  Middle %.3f'%self.ramp_middle,'ramp_middle'],
            [ 3, '  High %.3f'%self.ramp_high,'ramp_high'],
            [ 2, 'Exit Electrostatics Wizard','cmd.set_wizard()'],
            ]
        return result

    def get_prompt(self):
        self.prompt = None
        if self.showHelpPrompt:
            self.prompt = self.getHelpPrompt()
        return self.prompt

    def hide_help(self):
        self.showHelpPrompt = False
        cmd.refresh_wizard()

    def show_help(self):
        self.showHelpPrompt = True
        cmd.refresh_wizard()

    def isosurf(self,val,name='_is'):
        """Convenience routine for showing isosurfaces"""
        if val is None:
            cmd.delete(name)
        else:
            cmd.do('isosurf %s, %s, %.2f'%(name,self.map,val))
            if val < 0:
                cmd.do('color red, %s'%name)
                self.iso_neg = val
            else:
                cmd.do('color blue, %s'%name)
                self.iso_pos = val
        cmd.refresh_wizard()

    def surface(self,val,selection=None):
        """Convenience routine for showing molecular surfaces

        val is the amt by which we increase atom.vdw before drawing the surface.
        val==None means hide the surface.
        Doing this by tweaking the vdw of each atom seems like *REALLY* the wrong
        way to go.
        """
        if selection is None:
            selection = '(%s)'%self.molecule
        #
        # Reset to original vdw's
        #
        if self.vdw_modifier is not None:
            cmd.do('alter %s,vdw=str(float(vdw)-%s)'%(selection,self.vdw_modifier))
        #
        # Add our current modifier
        #
        self.vdw_modifier = val
        if self.vdw_modifier is not None:
            cmd.do('alter %s,vdw=str(float(vdw)+%s)'%(selection,self.vdw_modifier))
        # after an alter, a sort is required
        cmd.sort()
        if self.vdw_modifier is not None:
            cmd.show('surface',selection)
            self.ramp()
        if self.vdw_modifier is None:
            cmd.hide('surface',self.molecule)
            #cmd.hide('everything','e_lvl') # doesn't work :(
        cmd.refresh_wizard()
        
    def ramp(self,low=None,middle=None,high=None,
             ramp_name='e_lvl',map_name=None,molecule_name=None):
        """Convenience routine for setting up surface potential ramp

        If I were smart, this wouldn't make a new ramp each time ...
        """
        if map_name is None:
            map_name = self.map
        if molecule_name is None:
            molecule_name = self.molecule
        #
        # Callers can just specify one of (low,middle,high).
        #
        if low is not None:
            self.ramp_low = low
        if middle is not None:
            self.ramp_middle = middle
        if high is not None:
            self.ramp_high = high
        cmd.delete('ramp_name')
        cmd.ramp_new(ramp_name,map_name,[self.ramp_low,self.ramp_middle,self.ramp_high])
        print "new ramp",ramp_name,map_name,[self.ramp_low,self.ramp_middle,self.ramp_high]
        cmd.set('surface_color',ramp_name,molecule_name)
        print "set surface_color",ramp_name,molecule_name
        cmd.show('surface',molecule_name)
        cmd.sort()
        cmd.refresh_wizard()