import numpy
import sys
Method = sys.argv[1] # '3CNN' or 'Voxel_SVM'

all_pos_prob = numpy.load("../results/prob_scores/"+Method+"/prob_y.dat")
all_y_true = numpy.load("../results/prob_scores/"+Method+"/true_y.dat")

if Method == 'Voxel_SVM':
	all_pos_prob = all_pos_prob[:,1]

pos_ptf_file = open('../data/ptf/test_pos.ptf')
neg_ptf_file = open('../data/ptf/test_neg.ptf')

pos_ptf_list = list(pos_ptf_file)
neg_ptf_list = list(neg_ptf_file)

combined_list = pos_ptf_list + neg_ptf_list

PDB_dict = {}

for i in range(all_pos_prob.shape[0]):
	prob = all_pos_prob[i]
	line = combined_list[i]
	ele = line.split()
	PDB_ID = ele[0]
	x_ = float(ele[1])
	y_ = float(ele[2])
	z_ = float(ele[3])
	chain = ele[5]
	res = ele[-2]
	res_no = int(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:
	print PDB_ID
	pdb_file = open('../data/ATP_PDB/test/'+PDB_ID.lower()+'.pdb')
	out_file = open('../results/test_detect/'+PDB_ID.lower()+'_'+Method+'.pdb','w')

	model_ID = 0
	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')

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