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

default_mode = 'labchg'

class Charge(Wizard):

    def __init__(self,_self=cmd):
        Wizard.__init__(self,_self)

        self.modes = [
            'labchg',
            'addchg',
            'cpychg',
            'zrochg',
            'mzochg',
            'cbachg',
            'movchg',
            'rbachg',
            'sumchg',
            ]

        self.mode = default_mode
        self.status = 0

        self.mode_name = {
            'labchg':'Show',
            'cpychg':'Copy',
            'addchg':'Add',
            'zrochg':'Zero',
            'mzochg':'Move & Zero (atom)',
            'cbachg':'Move & Zero (resi)',
            'movchg':'Move & Remove (atom)',
            'rbachg':'Move & Remove (resi)',
            'sumchg':'Get Total Charge'
            }

        # initialize mode menu

        smm = []
        smm.append([ 2, 'Atom Charge Mode', '' ])
        for a in self.modes:
            smm.append([ 1, self.mode_name[a], 'cmd.get_wizard().set_mode("'+a+'")'])

        self.menu['mode']=smm

        self.memory = 0

        _self.edit_mode()


    def get_panel(self):
        return [
            [ 1, 'Charge Wizard',''],
            [ 3, self.mode_name[self.mode],'mode'],
            [ 2, 'Clear','cmd.get_wizard().clear()'],
            [ 2, 'Done','cmd.set_wizard()'],
            ]

    def cleanup(self):
        global default_mode
        default_mode = self.mode
        self.clear()

    def clear(self):
        self.set_status(0)
        if 'wcharge' in self.cmd.get_names('selections'):
            if self.mode!='sumchg':
                self.cmd.edit("wcharge")
                self.cmd.label("pkmol",'') # fastest clear command
            else:
                self.cmd.label("wcharge",'') # fastest clear command
            self.cmd.delete("wcharge")
            self.cmd.unpick()
        self.cmd.unpick()
        self.cmd.refresh_wizard()

    def get_prompt(self):
        self.prompt = None
        if self.mode == 'cpychg':
            if self.status==0:
                self.prompt = [ 'Pick source atom...' ]
            elif self.status==1:
                self.prompt = [ 'Pick destination atom on which to assign charge %6.4f'%self.partial_charge ]
        if self.mode == 'addchg':
            if self.status==0:
                self.prompt = [ 'Pick source atom...' ]
            elif self.status==1:
                self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ]
        if self.mode == 'mzochg':
            if self.status==0:
                self.prompt = [ 'Pick source atom to copy and zero...' ]
            elif self.status==1:
                self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ]
        if self.mode == 'movchg':
            if self.status==0:
                self.prompt = [ 'Pick source atom to copy and destroy...' ]
            elif self.status==1:
                self.prompt = [ 'Pick destination atom on which to add charge %6.4f'%self.partial_charge ]
        if self.mode == 'sumchg':
            if self.status==0:
                self.prompt = [ 'Pick an atom on the chain...' ]
            if self.status==1:
                self.prompt = [ 'Total charge on the chain is %6.4f'%self.partial_charge,
                                     'Pick an atom on the chain...' ]
        if self.mode == 'cbachg':
            if self.status==0:
                self.prompt = [ 'Pick source residue to copy and zero...' ]
            elif self.status==1:
                self.prompt = [ 'Pick destination residue on which to add charges.']
        if self.mode == 'rbachg':
            if self.status==0:
                self.prompt = [ 'Pick source residue to copy and remove...' ]
            elif self.status==1:
                self.prompt = [ 'Pick destination residue on which to add charges.']
        if self.mode == 'zrochg':
                self.prompt = [ 'Pick atom on which to zero charge...' ]


        if self.mode == 'labchg':
                self.prompt = [ 'Pick atom on which to show charge...' ]

        if "wcharge" in self.cmd.get_names('selections'):
            pymol.stored.charge = 0
            if self.cmd.iterate("(byres wcharge)",
                                "stored.charge = stored.charge + partial_charge"):
                self.prompt.insert(0,"Total charge on the residue is %6.4f"%pymol.stored.charge)

        return self.prompt

    def set_mode(self,mode):
        if mode in self.modes:
            self.mode = mode
        self.status = 0
        self.cmd.refresh_wizard()

    def set_status(self,status):
        self.status = status
        self.cmd.refresh_wizard()

    def do_pick(self,bondFlag):
        if bondFlag:
            print(" Error: please select a single atom")

        if self.mode == 'cpychg':
            # picking up
            if self.status==0:
                if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
                    self.partial_charge = pymol.stored.charge
                    self.status = 1
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

            # dropping off
            elif self.status==1:
                pymol.stored.charge=self.partial_charge
                if self.cmd.alter("(pk1)","partial_charge = stored.charge"):
                    self.status = 0
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

        if self.mode == 'cbachg' or self.mode == 'rbachg':
            # picking up
            if self.status==0:
                pymol.stored.chg_dict = {}
                if self.cmd.iterate("(byres pk1)","stored.chg_dict[name] = partial_charge"):
                    self.charge_dict = pymol.stored.chg_dict
                    self.status = 1
                    self.cmd.label("(byres pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(byres pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

            # dropping off
            elif self.status==1:
                pymol.stored.valid_atoms = []
                if self.cmd.iterate("(byres pk1)","stored.valid_atoms.append(name)"):
                    kees = list(self.charge_dict.keys())
                    valid_dict = {}
                    for a in pymol.stored.valid_atoms:
                        if a in kees:
                            valid_dict[a] = 1
                    pymol.stored.chg_dict = self.charge_dict
                    # copy/add charges
                    for a in valid_dict.keys():
                        self.cmd.alter("((byres pk1) and name %s)"%a,
                                     "partial_charge = partial_charge + stored.chg_dict[name]")
                        if self.mode == 'rbachg':
                            self.cmd.remove("((wcharge) and name %s)"%a)
                        else:
                            self.cmd.alter("((wcharge) and name %s)"%a,
                                         "partial_charge = 0")
                    self.status = 0
                    # update labels
                    self.cmd.label("(wcharge or (byres pk1))","'%6.4f'%partial_charge")
                    # show which atoms had charges moved
                    self.cmd.select("wcharge","(none)")
                    for a in valid_dict.keys():
                        self.cmd.select("wcharge","(wcharge or ((byres pk1) and name %s))"%a)
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

        if self.mode == 'addchg':
            # picking up
            if self.status==0:
                if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
                    self.partial_charge = pymol.stored.charge
                    self.status = 1
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

            # dropping off
            elif self.status==1:
                pymol.stored.charge=self.partial_charge
                if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"):
                    self.status = 0
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

        if self.mode == 'mzochg':
            # picking up
            if self.status==0:
                if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
                    self.partial_charge = pymol.stored.charge
                    self.status = 1
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

            # dropping off
            elif self.status==1:
                pymol.stored.charge=self.partial_charge
                if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"):
                    self.status = 0
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.alter("(wcharge)","partial_charge=0")
                    self.cmd.label("(wcharge)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

        if self.mode == 'movchg':
            # picking up
            if self.status==0:
                if self.cmd.iterate("(pk1)","stored.charge = partial_charge"):
                    self.partial_charge = pymol.stored.charge
                    self.status = 1
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

            # dropping off
            elif self.status==1:
                pymol.stored.charge=self.partial_charge
                if self.cmd.alter("(pk1)","partial_charge = partial_charge + stored.charge"):
                    self.status = 0
                    self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                    self.cmd.remove("wcharge")
                    self.cmd.select("wcharge","(pk1)")
                    self.cmd.unpick()
                    self.cmd.enable("wcharge")

        if self.mode == 'zrochg':
            if self.cmd.alter("(pk1)","partial_charge = 0.0"):
                self.cmd.label("(pk1)","'%6.4f'%partial_charge")
                self.cmd.select("wcharge","(pk1)")
                self.cmd.unpick()
                self.cmd.enable("wcharge")

        if self.mode == 'labchg':
            self.cmd.label("(pk1)","'%6.4f'%partial_charge")
            self.cmd.select("wcharge","(pk1)")
            self.cmd.unpick()
            self.cmd.enable("wcharge")

        if self.mode == 'sumchg':
            pymol.stored.charge = 0.0
            if self.cmd.iterate("(pkmol)","stored.charge = stored.charge + partial_charge"):
                self.partial_charge = pymol.stored.charge
                self.status = 1
                self.cmd.select("wcharge","(pkmol)")
                self.cmd.unpick()
                self.cmd.enable("wcharge")

        self.cmd.refresh_wizard()
