#
# autoMultiFit.py -- Given a homo-oligomer in N-states, fit any combination of chains across states
#
# AUTHOR: Jason Vertrees
# DATE  : 2009-11-02
#
import pymol
from pymol import cmd
 
def autoMultiFit(sel, fromChain="*", fromState="*", toChain="*", toState="*", verbose=None):
    """
    FUNCTION
      Given a homo-oligomer in N-states (say from MD), fit any
      combination of chains/selections across states. See Usage, Examples and Notes.
 
    USAGE
      autoMultiFit sel, [fromChain[, fromState[, toChain[, toState[, verbose ]]]]]
 
      sel
        The selection that must exist in all chains (what you're fitting)
 
      fromChain
        Fit from this chain, leave blank or use '*' for all chains;
 
      fromState
        Fit from this state, leave blank or use '*' for all states;
 
      toChain
        Fit to this chain only, use '*' or leave blank for all chains;
 
      toState
        Fit to this state only, use '*' or leave blank for all states
 
      verbose
        If turned on, this prints the fit for all pairs; this could be VERY slow.
 
    RETURNS
      A hash of hashes of hashes of hashes. :-)  Access the results like this:
 
      fitValues = autoMultiFit( mySelection )
      #fitValues[fromChain][fromState][toChain][toState]
      # to get the result of the fit from chain A, state 4 to chain C, state 17 type,
      fitValues['A']['4']['C']['17']
 
    EXAMPLES
      * Fit ALL chains to each other across ALL STATES for PDB 2kb1
      autoMultiFit 2kb1
 
      * Fit chain A's residue 22-34 against all states
      autoMutliFit 2kb1 and i. 22-34, fromChain=A
 
      * Fit chain A, B and C's residues 8-17, states 10 and 11, against chain D's state 8
      myVals = autoMultiFit("2kb1 and i. 8-17", fromChain="a b c", fromState="10 11", toChain="d", toState=8)
      myVals['B']['10']['D']['8']
 
      NOTES
        * The return value uses UPPERCASE and STRINGS for ALL hash keys, so if you used chain 'a' it will
          be 'A' in the map; if you used 'chainXX' it will be 'CHAINXX' in the map.
        * Continuing from the above note, all statese are accessible by their strings, so '17' gets state
          17 from the hash, not just the number 17.
        * This can be very slow: it's O( nFromChains * nFromStates * nToChains * nToStates )
        * fromChain, toChain, fromStates and toStates can all be single values, like "8", they can be a SPACE
          separated list of values, like "2 4 5", or they can be "*" which means ALL.
    """
 
    fromChain = processChain(fromChain,sel)
    fromState = processState(fromState, sel)
    toChain   = processChain(toChain, sel)
    toState   = processState(toState, sel)
 
    res = {}
 
    for frC in fromChain:
        res[frC] = {}
        for frS in fromState:
            res[frC][str(frS)] = {}
            cmd.create( "__tmpA", sel + " and c. " + frC, frS, 1 )
            for toC in toChain:
                res[frC][str(frS)][toC] = {}
                for toS in toState:
                    cmd.create("__tmpB", sel + " and c. " + toC, toS, 1 )
                    curVal = cmd.pair_fit( "__tmpA", "__tmpB", quiet=1 )
                    res[frC][str(frS)][toC][str(toS)] = curVal
                    if verbose!=None:
                        print "Pair Fitted: from (chain: %s, state: %s) to (chain: %s, states %s) with RMSD %s" % (frC, frS, toC, toS, curVal)
                    cmd.delete("__tmpB")
            cmd.delete("__tmpA")
    return res
 
def processChain(ch, sel):
    """Turn a string-based space separated list into an array of chains"""
    if ch == "*":
        # just in case
        return map(str.upper, cmd.get_chains(sel))
    else:
        return map(str.upper, processTextList(ch))
 
def processState(st, sel):
    """Tur a string-based space separated list into an array of states"""
    if st == "*":
        # assumes that if there is a state X, there is a state, X-1
        return range(cmd.count_states(sel))
    else:
        return map(int, processTextList(st))
 
def processTextList(t):
    """ make a space-separated list into a real Python list, eg.
    input: a b c d
    output: [ 'a', 'b', 'c', 'd' ]
    """
    t = str(t).split()
    return t
 
 
if __name__ == "__main__":
    assert processTextList( "a b c d" ) == ['a', 'b', 'c', 'd']
    assert processTextList( " 1 45 s s s s j p k") == [ '1', '45', 's', 's', 's', 's', 'j', 'p', 'k' ]
 
    assert processChain( "a b c dd", "blankSel" ) == ["a", "b", "c", "dd"]
    assert processState( "1 2 3 44 5", "blankSel" ) == [ 1, 2, 3, 44, 5 ]
 
 
cmd.extend( "autoMultiFit", autoMultiFit )