from pymol import cmd
file=open("C:/Users/vladm/Desktop/pymol_scripts/contacts.txt","w")
cmd.fetch("1ozj")
cmd.create("dna","(/1ozj/A/C)+(/1ozj/B/D)")
cmd.create("protein","(/1ozj/C/A)+(/1ozj/D/B)")
def atom_search(resn,atom):
	return("/dna/*/*/"+resn+"/"+atom)
data={"deoxyrib_polar":["*",["O*'"]],
"deoxyrib_nonpolar":["*",["C*'"]],
"phosphate_polar":["*",["OP*"]],
"phosphate_nonpolar":["*",["P"]],
"A_big_polar":["DA",["N3","N9"]],
"A_small_polar":["DA",["N1","N6","N7"]],
"G_big_polar":["DG",["O6","N7"]],
"G_small_polar":["DG",["N2","N3","N9"]],
"T_big_polar":["DT",["N1","O2"]],
"T_small_polar":["DT",["O4"]],
"C_big_polar":["DC",["N1","O2","N3"]],
"C_small_polar":["DC",["N4"]],
"A_big_nonpolar":["DA",["C2","C4","C8"]],
"A_small_nonpolar":["DA",["C5","C6"]],
"G_big_nonpolar":["DG",["C5","C6","C8"]],
"G_small_nonpolar":["DG",["C2","C4"]],
"T_big_nonpolar":["DT",["C2","O2"]],
"T_small_nonpolar":["DT",["C4","C5","C7"]],
"C_big_nonpolar":["DC",["C2"]],
"C_small_nonpolar":["DC",["C4","C5","C6"]]}
for contact_type in data.keys():
	count=[]
	file.write(contact_type+"\n")
	if "non" in contact_type:
		for elem in data[contact_type][1]:
			cmd.iterate("((protein/*/*/*/C*)+(protein/*/*/*/S*)) within 4.5 of "+atom_search(data[contact_type][0],elem),"count.append('1')")
	else:
		for elem in data[contact_type][1]:
			cmd.iterate("((protein/*/*/*/N*)+(protein/*/*/*/O*)) within 3.5 of "+atom_search(data[contact_type][0],elem),"count.append('1')")
	file.write(str(len(count))+"\n")
cmd.delete("all")
file.close()