# Warren L. DeLano

from pymol.wizard import Wizard
from pymol import cmd
import pymol

saved_mode = 2 # toggle
saved_scope = 1 # byres
saved_what = 3 # sticks
saved_color = 5 # magenta

class Appearance(Wizard):


    mode_dict = {
        0 : ['Color', '_ cmd.color' ],
        1 : ['Color (elem c)', '_ util.color_carbon' ],
        2 : ['Toggle','_ cmd.toggle'],
        3 : ['Show', '_ cmd.show'],
        4 : ['Hide', '_ cmd.hide'],
        5 : ['Select', '_ cmd.select'],
        }

    what_dict = {
        1 : ['Lines', 'lines'],
        2 : ['Nonbonded', 'nonbonded' ],
        3 : ['Sticks', 'sticks' ],
        4 : ['Ribbon', 'ribbon' ],
        5 : ['Cartoon', 'cartoon' ],
        6 : ['Labels', 'labels' ],
        7 : ['Dots', 'dots' ],
        8 : ['Spheres', 'spheres' ],
        9 : ['NB Spheres', 'nb_spheres' ],
        10: ['Mesh', 'mesh' ],
        11: ['Surface', 'surface' ],
        }

    color_dict = {
        1 : ['\\900red', 'red' ],
        2 : ['\\090green', 'green' ],
        3 : ['\\009blue'  , 'blue' ],
        4 : ['\\990yellow' , 'yellow' ],
        5 : ['\\909magenta', 'magenta' ],
        6 : ['\\099cyan'    ,  'cyan' ],
        7 : ['\\955salmon'   , 'salmon' ],
        8 : ['\\595lime'      , 'lime' ],
        9 : ['\\967pink'  , 'pink' ],
        10 : ['\\559slate' ,  'slate' ],
        11 : ['\\949violet' , 'violet' ],
        12 : ['\\950orange'  , 'orange' ],
        13 : ['\\059marine'   ,  'marine' ],
        14 : ['\\905hotpink' , 'hotpink' ],
        }

    scope_dict = {
        0: ['By Atom', ''],
        1: ['By Residue', 'byres' ],
        2: ['By Chain', 'bychain' ],
        3: ['By Segment', 'bysegment' ],
        4: ['By Object', 'byobject' ],
        5: ['By Molecule', 'bymol' ],
        }

    def __init__(self,_self=cmd):
        Wizard.__init__(self,_self)
        self.cmd.deselect()
        self.cmd.unpick()
        self.selection_mode = self.cmd.get_setting_int("mouse_selection_mode")
        self.cmd.set("mouse_selection_mode",0) # set selection mode to atomic
        self.current_mode = saved_mode
        self.current_what = saved_what
        self.current_scope = saved_scope
        self.current_color = saved_color
        self.menu['mode'] = [
            [2, 'Mode', '' ],
            [1, self.mode_dict[0][0], 'cmd.get_wizard().set_mode(0)' ],
            [1, self.mode_dict[1][0], 'cmd.get_wizard().set_mode(1)' ],
            [ 0, ''           , '' ],
            [1, self.mode_dict[2][0], 'cmd.get_wizard().set_mode(2)' ],
            [1, self.mode_dict[3][0], 'cmd.get_wizard().set_mode(3)' ],
            [1, self.mode_dict[4][0], 'cmd.get_wizard().set_mode(4)' ],
#         [1, self.mode_dict[5][0], 'cmd.get_wizard().set_mode(4)' ],
            ]

        self.menu['what'] = [
            [2, 'What', '' ],
            [1, self.what_dict[1][0], 'cmd.get_wizard().set_what(1)' ],
            [1, self.what_dict[2][0], 'cmd.get_wizard().set_what(2)' ],
            [1, self.what_dict[3][0], 'cmd.get_wizard().set_what(3)' ],
            [1, self.what_dict[4][0], 'cmd.get_wizard().set_what(4)' ],
            [1, self.what_dict[5][0], 'cmd.get_wizard().set_what(5)' ],
            [0, '', '' ],
            [1, self.what_dict[6][0], 'cmd.get_wizard().set_what(6)' ],
            [0, '', '' ],
            [1, self.what_dict[7][0], 'cmd.get_wizard().set_what(7)' ],
            [1, self.what_dict[8][0], 'cmd.get_wizard().set_what(8)' ],
            [1, self.what_dict[9][0], 'cmd.get_wizard().set_what(9)' ],
            [0, '', '' ],
            [1, self.what_dict[10][0], 'cmd.get_wizard().set_what(10)' ],
            [1, self.what_dict[11][0], 'cmd.get_wizard().set_what(11)' ],
            ]

        self.menu['color'] = [
            [2, 'Color', '' ],
            [1, self.color_dict[1][0], 'cmd.get_wizard().set_color(1)' ],
            [1, self.color_dict[2][0], 'cmd.get_wizard().set_color(2)' ],
            [1, self.color_dict[3][0], 'cmd.get_wizard().set_color(3)' ],
            [1, self.color_dict[4][0], 'cmd.get_wizard().set_color(4)' ],
            [1, self.color_dict[5][0], 'cmd.get_wizard().set_color(5)' ],
            [1, self.color_dict[6][0], 'cmd.get_wizard().set_color(6)' ],
            [1, self.color_dict[7][0], 'cmd.get_wizard().set_color(7)' ],
            [1, self.color_dict[8][0], 'cmd.get_wizard().set_color(8)' ],
            [1, self.color_dict[9][0], 'cmd.get_wizard().set_color(9)' ],
            [1, self.color_dict[10][0], 'cmd.get_wizard().set_color(10)' ],
            [1, self.color_dict[11][0], 'cmd.get_wizard().set_color(11)' ],
            [1, self.color_dict[12][0], 'cmd.get_wizard().set_color(12)' ],
            [1, self.color_dict[13][0], 'cmd.get_wizard().set_color(13)' ],
            [1, self.color_dict[14][0], 'cmd.get_wizard().set_color(14)' ],
            ]

        self.menu['scope'] = [
            [ 2, 'Scope', ''],
            [ 1, self.scope_dict[0][0], 'cmd.get_wizard().set_scope(0)' ],
            [ 1, self.scope_dict[1][0], 'cmd.get_wizard().set_scope(1)' ],
            [ 1, self.scope_dict[2][0], 'cmd.get_wizard().set_scope(2)' ],
            [ 1, self.scope_dict[3][0], 'cmd.get_wizard().set_scope(3)' ],
            [ 1, self.scope_dict[4][0], 'cmd.get_wizard().set_scope(4)' ],
            [ 0, ''           , '' ],
            [ 1, self.scope_dict[5][0], 'cmd.get_wizard().set_scope(5)' ],
            ]

    def set_scope(self,scope):
        scope = int(scope)
        if scope in self.scope_dict:
            self.current_scope = scope
        self.cmd.refresh_wizard()

    def set_what(self,what):
        what = int(what)
        if what in self.what_dict:
            self.current_what = what
        self.cmd.refresh_wizard()

    def set_color(self,color):
        color = int(color)
        if color in self.color_dict:
            self.current_color = color
        self.cmd.refresh_wizard()

    def set_mode(self,mode):
        mode = int(mode)
        if mode in self.mode_dict:
            self.current_mode = mode
        self.cmd.refresh_wizard()

    def undo(self):
        print("no undo!")

    def get_prompt(self):
        self.prompt = []
        return self.prompt

    def get_panel(self):
        panel = [
                [ 1, 'Appearance Wizard',''],
                [ 3, self.mode_dict[self.current_mode][0], 'mode' ],
                ]
        if self.current_mode in [0,1]: # color mode
            panel.append(
                [ 3, self.color_dict[self.current_color][0], 'color' ])
        elif self.current_mode in [2,3,4]: # what mode
            panel.append(
            [ 3, self.what_dict[self.current_what][0], 'what' ])
        else: # select mode
            panel.append(
                [ 1, 'Atoms' , '' ])

        panel.extend([
                [ 3, self.scope_dict[self.current_scope][0], 'scope' ],
                #         [ 2, 'Undo', 'cmd.get_wizard().undo()' ],
                [ 2, 'Done','cmd.set_wizard()'],
                ])
        return panel


    def do_pick(self,bondFlag):
        if self.current_mode in [0,1]: # color
            sele = "(%s pk1)"%self.scope_dict[self.current_scope][1]
            color = self.color_dict[self.current_color][1]
            mode = self.mode_dict[self.current_mode][1]
            cmmd = mode+'("%s","%s")'%(color,sele)
            self.cmd.do(cmmd,log=0)
        elif self.current_mode in [2,3,4]: # show/hide/toggle
            sele = "(%s pk1)"%self.scope_dict[self.current_scope][1]
            what = self.what_dict[self.current_what][1]
            mode = self.mode_dict[self.current_mode][1]
            cmmd = mode+'("%s","%s")'%(what,sele)
            self.cmd.do(cmmd,log=0)
        else: # select
            pass
        self.cmd.unpick()
        self.cmd.refresh_wizard()
        return 1

    def do_select(self,selection):
        if self.current_mode in [0,1]: # color
            sele = "(%s %s)"%(self.scope_dict[self.current_scope][1],selection)
            color = self.color_dict[self.current_color][1]
            mode = self.mode_dict[self.current_mode][1]
            cmmd = mode+'("%s","%s")'%(color,sele)
            self.cmd.do(cmmd,log=0)
        elif self.current_mode in [2,3,4]: # show/hide/toggle
            sele = "(%s %s)"%(self.scope_dict[self.current_scope][1],selection)
            what = self.what_dict[self.current_what][1]
            mode = self.mode_dict[self.current_mode][1]
            cmmd = mode+'("%s","%s")'%(what,sele)
            self.cmd.do(cmmd,log=0)
        self.cmd.delete(selection)
        self.cmd.deselect()
        self.cmd.unpick()
        self.cmd.refresh_wizard()
        return 1

    def cleanup(self):
        self.cmd.set("mouse_selection_mode",self.selection_mode) # restore selection mode
        global saved_scope
        saved_scope = self.current_scope
        global saved_mode
        saved_mode = self.current_mode
        global saved_color
        saved_color = self.current_color
        global saved_what
        saved_what = self.current_what
