# -----------------------------------------------------------------------------
# Trace C-alpha backbone adding one residue at a time and show simulated map.
#

pdb_path = '/usr/local/src/staff/goddard/pdb/pdb1fav.ent'
resolution = 5.0                # Simulated density resolution, Angstroms
map_color = (.6,.7,.8,.3)        # red,green,blue,opacity.

from chimera import runCommand

# Open model
from chimera import openModels as om
m = om.open(pdb_path, 'PDB')[0]

# Orient molecule
runCommand('turn y 90')
runCommand('turn x 30')

# Color molecule and background
runCommand('rainbow')
runCommand('set bg_color white')

# Create residue number label
runCommand('2dlabel create rlab xpos 0.1 ypos 0.9')

# Set display style.
runCommand('repr stick #%d' % m.id)
m.stickScale = 5.0      # Fatter sticks

# Order residues by number and chain id.
rlist = m.residues
rlist.sort(key = lambda r: (r.id.position, r.id.chainId))

# Hide all atoms.
for a in m.atoms:
    a.display = False

# Show density adding one residue at a time.
from MoleculeMap import molecule_map
atoms = []
map = None
for r in rlist:
    # Accumulate atoms from residue.
    atoms.extend(r.atoms)

    # Close previous map.
    if map:
        om.close([map])

    # Make map.
    map = molecule_map(atoms, resolution, gridSpacing = .2*resolution,
                       showDialog = False)
    map.set_parameters(surface_colors = [map_color],
                       one_transparent_layer = True)
    map.show()

    # Display backbone atoms.
    for a in r.atoms:
        if a.name == 'CA':
            a.display = True

    # Update residue label
    text = 'chain %s residue %d' % (r.id.chainId, r.id.position)
    if r.isHelix:
        text += ' helix'
    elif r.isSheet:
        text += ' sheet'
    runCommand('2dlabel change rlab text "%s" color black' % text)

    # Draw one frame.
    runCommand('wait 1')
