import numpy
import sys


def visualize_detected_sites(numpy_ID_list,site_ID):
	PDB_dict = {}
	for numpy_ID in numpy_ID_list:
		summary_file = open('../results/summary/'+'PDB_max_fold_prob_'+numpy_ID+'.txt')
		summary_list = list(summary_file)
		for line in summary_list:
			ele = line.split()
			PDB_ID = ele[0]
			x_ = float(ele[2])
			y_ = float(ele[3])
			z_ = float(ele[4])
			res = ele[5]
			chain = ele[6]
			res_no = int(ele[7])
			prob = float(ele[-1])

			if PDB_ID not in PDB_dict:
				PDB_dict[PDB_ID]={}
			PDB_dict[PDB_ID][(res,chain,res_no)]=prob

	for PDB_ID in PDB_dict:
		pdb_file = open('../data/pdb/'+PDB_ID.lower()+'.pdb')
		out_file = open('../results/pdb_detect/'+PDB_ID.lower()+'_detect_'+site_ID+'.pdb','w')
		for line1 in pdb_file:
			line=line1.split()
			if line[0]=="ATOM": 
				atom=(line1[13:16].strip(' '))
				res=(line1[17:20])
				chain_ID=line1[21:26]
				chain=chain_ID[0]
				res_no=chain_ID[1:].strip(' ')
				res_no=int(res_no)
				chain_ID=(chain,res_no)
				new_pos=[float(line1[30:37]),float(line1[38:45]),float(line1[46:53])]
				if (res,chain,res_no) in PDB_dict[PDB_ID].keys():
					out_file.write(line1[0:60]+'\t'+str("%.2f" % PDB_dict[PDB_ID][(res,chain,res_no)]).rjust(6)+'           '+str(atom[0])+'\n')
				else:
					out_file.write(line1[0:60]+'\t'+str("%.2f" % 0).rjust(6)+'           '+str(atom[0])+'\n')

			if line[0]=="ENDMDL":
				break