'''
(c) 2010 Thomas Holder, MPI for Developmental Biology
'''

from pymol import cmd, stored
from pymol.exporting import _resn_to_aa as one_letter
import pymol.wizard.mutagenesis as mutagenesis
import itertools

three_letter = dict(reversed(i) for i in one_letter.iteritems())

def mutate(selection, new_resn, inplace=0, sculpt=0):
    '''
DESCRIPTION

    Mutate a single residue. Does call the mutagenesis wizard non-interactively
    and tries to select the best rotamer. Can do some sculpting in the end to
    the best rotamer.

USAGE

    mutate selection, new_resn [, inplace [, sculpt ]]

ARGUMENTS

    selection = string: atom selection of single residue

    new_resn = string: new residue name (3-letter or 1-letter)

    inplace = 0 or 1: work on copy of input object if 0 {default: 0}

    sculpt = 0: no sculpting {default}

    sculpt = 1: do sculpting on best rotamer

    sculpt = 2: do sculpting with neighbors

EXAMPLE

    fetch 2x19, async=0
    select x, A/CYS`122/
    zoom x
    mutate x, LYS
    '''
    inplace, sculpt = int(inplace), int(sculpt)
    org = cmd.get_object_list(selection)[0]
    tmp = cmd.get_unused_name()
    new_resn = new_resn.upper()
    new_resn = three_letter.get(new_resn, new_resn)

    if inplace:
        cpy = org
    else:
        cpy = cmd.get_unused_name(org + '_cpy')
        cmd.create(cpy, org, -1, 1, zoom=0)

    cmd.iterate('first (%s)' % selection, 'stored.x = (segi,chain,resi)')
    res = '/%s/%s/%s/%s' % ((cpy,) + stored.x)
    cmd.remove('%s and not name CA+CB+C+N+O+OXT' % (res))

    # start the wizard to count the number of rotamers for this residue
    cmd.wizard("mutagenesis")
    cmd.get_wizard().set_mode(new_resn)
    cmd.get_wizard().do_select("("+res+")")

    # foreach state, make the change and write to disk
    best_state = (1, 1e9)
    for i in range(1, cmd.count_states(mutagenesis.obj_name)+1):
        cmd.create(tmp, mutagenesis.obj_name + ' and not name CA+CB+C+N+O', i, 1, zoom=0)
        c4 = cmd.count_atoms(cpy + ' within 4 of ' + tmp, state=1)
        c3 = cmd.count_atoms(cpy + ' within 3 of ' + tmp, state=1)
        c2 = cmd.count_atoms(cpy + ' within 2 of ' + tmp, state=1)
        score = 100*c2 + 10*c3 + 1*c4
        print 'Frame %d Score %.2f (c4: %d, c3: %d, c2: %d)' % (i, score, c4, c3, c2)
        if score < best_state[1]:
            best_state = (i, score)

    cmd.delete(tmp)
    print 'Best: Frame %d Score %.2f' % best_state
    cmd.frame(best_state[0])
    cmd.get_wizard().apply()
    cmd.wizard()

    if sculpt > 0:
        # do 100 iterations, 75 of them with all terms but low VDW weights,
        # and 25 with only local geometry terms. With default VDW weights and
        # atom clashes, the structure gets distorted very easily!
        cmd.protect('(not %s) or name CA+C+N+O+OXT' % (res))
        cmd.sculpt_activate(cpy)
        cmd.set('sculpt_vdw_weight', 0.25, cpy) # Low VDW forces
        cmd.set('sculpt_field_mask', 0x1FF, cpy) # Default
        if sculpt == 2:
            cmd.sculpt_iterate(cpy, cycles=25)
            cmd.deprotect('(byres (%s within 6.0 of %s)) and (not name CA+C+N+O+OXT)' % (cpy, res))
            cmd.sculpt_iterate(cpy, cycles=50)
        else:
            cmd.sculpt_iterate(cpy, cycles=75)
        cmd.set('sculpt_field_mask', 0x01F, cpy) # Local Geometry Only
        cmd.sculpt_iterate(cpy, cycles=25)
        cmd.unset('sculpt_vdw_weight', cpy)
        cmd.unset('sculpt_field_mask', cpy)
        cmd.sculpt_deactivate(cpy)
        cmd.deprotect()

cmd.extend('mutate', mutate)

# vi: expandtab:smarttab