# CM; label residues in PyMOL; June 2003
# import the basic Python glob module and PyMOL commands, etc.:
import sys
from glob import glob
from pymol import cmd

"""
AUTHOR:
	Cameron Mura (cmura@ucsd.edu)
SYNTAX:
	load module into PyMOL:	
		run label_residues.py
	run as: 	
		label_residues(selection,start_res,end_res,increment)

"""
def label_residues(selection="",start_res=0,end_res=0,increment=0):

	### check some chingis first:
	if selection == "":
		print "*** WARNING: NO 'selection' was defined -- using everything!\n";
		selection = "all";
	else:
		print "\n***********************************************";
		print "SELECTION that will be labeled: ",selection;
		print "***********************************************\n";

	if start_res == 0:
		print "*** WARNING: NO 'start_res' was defined -- using default of 1!\n";
		start_res = 1;
	else:
		start_res=int(start_res);
	if end_res == 0:
		print "*** WARNING: NO 'end_res' was defined -- using default of 50!\n";
		end_res = 50;
	else:
		end_res=int(end_res);
	if increment == 0:
		print "*** WARNING: NO 'increment' was defined -- using default of 10!\n";
		increment = 10;
	else:
		increment=int(increment);
	###

	#count residues using two ways to count atoms (could also use this in above line)
	num_atoms = len(cmd.get_model(selection).atom);
	num_residues = cmd.count_atoms("name ca and "+selection);
	
	if num_atoms == 0:
		print "*** WARNING: NO ATOMS in the selection \"%s\"!!\n" % (selection);
	else:
		print "Number of residues in selection \"%s\":\t%i" % (selection,num_residues);
		print "(consisting of %s atoms)" % (num_atoms);
	
	counter = start_res;
	while counter <= end_res:
		print "Labeling residue: ",counter;
		label_selection = selection + " and resi " + str(counter);
		cmd.label(label_selection,"resi");
		counter = counter + increment;
# bind external functions as commands in the PyMOL session...so can do function-calling without parentheses...
cmd.extend("label_residues",label_residues)