import os
import sys
import numpy
import scipy.ndimage
from sets import Set
import json
import os
from scipy import spatial
import theano
from atom_res_dict import *
from generate_backbone_box_20A import *

def pocket_pts_to_backbone_box(pts_list,input_dir,pdb_dir,max_nodes_in_poc=50):
	pocket_valid=True
	all_pocket=[]
	dat_num=0
	for fn in pts_list:
		FV=[] 
		PDB_ID = fn[0:4]
		f = open(os.path.join(input_dir,fn))
		infile=list(f)
		for line in infile:
			S=line.split()
			if S!=[]:
				PDB_ID = S[0]
				label= int(S[-1])
				pos=[float(S[1]),float(S[2]),float(S[3])]
				pos=numpy.array(pos)
				X, valid = cut_box(PDB_ID, pos, pdb_dir)
				if valid:
					FV.append((X,pos,label))  
				else:
					pocket_valid=False
					return all_pocket, pocket_valid

		f.close()
		num_res_in_poc = len(FV)

		if num_res_in_poc > max_nodes_in_poc:
			return all_pocket, False
		if num_res_in_poc<max_nodes_in_poc:
			for i in range(len(FV),max_nodes_in_poc):
				 dummpy = numpy.zeros(X.shape)
				 FV.append((dummpy,numpy.array([0,0,0]),-1))

		vectors = FV
		all_pocket.append([fn,vectors,num_res_in_poc])

	return all_pocket, pocket_valid

def cut_box(PDB_ID,pos,pdb_dir = '../../3DCNN_data_backbone/data/PDB_family_train'):
	
	pixel_size = 1
	num_of_channels = 4
	std = 0.6
	valid = False
   
	total_num_of_samples=0
	box_size=20
	num_3d_pixel=int(box_size/pixel_size)
	box_x_min=-box_size/2
	box_x_max=+box_size/2
	box_y_min=-box_size/2
	box_y_max=+box_size/2
	box_z_min=-box_size/2
	box_z_max=+box_size/2

	x,y,z = pos[0],pos[1],pos[2]

	if not os.path.isfile(pdb_dir+'/'+PDB_ID.lower()+'.pdb'):
		return None, False

	pdb_file=open(pdb_dir+'/'+PDB_ID.lower()+'.pdb')


	l=list(pdb_file)
	PROTEIN = grab_PDB(l)
	[ID_dict, all_pos, all_lines, all_atom_type, PDB_entries, all_x, all_y , all_z] = PROTEIN

	all_pos = numpy.array(all_pos)
	my_kd_tree = scipy.spatial.KDTree(all_pos)
	actual_atom=find_actual_pos(my_kd_tree, [x,y,z], PDB_entries)
	chain_ID=actual_atom.chain_ID
	res_atoms=ID_dict[chain_ID]
	res=actual_atom.res

	sample=numpy.zeros((num_of_channels,box_size/pixel_size,box_size/pixel_size,box_size/pixel_size))
	X_smooth=numpy.zeros(sample.shape, dtype=theano.config.floatX)

	if res in res_label_dict.keys():
		label = res_label_dict[res]
		get_position=get_position_dict(res_atoms)
		if "CA" in Set(get_position.keys()) and "C" in Set(get_position.keys()) and "N" in Set(get_position.keys()):
			backbone=ID_dict[chain_ID][0:4]
			[reference,transform]=center_and_transform(label,get_position)   
			transformed_pos = ((all_pos - reference).dot(transform))-bias
			x_index = numpy.intersect1d(numpy.where(transformed_pos[:,0]>box_x_min),numpy.where(transformed_pos[:,0]<box_x_max))
			y_index = numpy.intersect1d(numpy.where(transformed_pos[:,1]>box_y_min),numpy.where(transformed_pos[:,1]<box_y_max))
			z_index = numpy.intersect1d(numpy.where(transformed_pos[:,2]>box_z_min),numpy.where(transformed_pos[:,2]<box_z_max))

			final_index = numpy.intersect1d(x_index,y_index)
			final_index = numpy.intersect1d(final_index,z_index)
			final_index = final_index.tolist()

			final_index = [ ind for ind in final_index if (all_atom_type[ind] =='N' or all_atom_type[ind]=='CA' or all_atom_type[ind]=='C' or all_atom_type[ind]=='O') ]

			box_ori = [PDB_entries[i] for i in final_index] 
			new_pos_in_box = transformed_pos[final_index]

			atom_count = len(box_ori)
	
			for i in range (0,len(box_ori)):
				atoms = box_ori[i]
				x=new_pos_in_box[i][0]
				y=new_pos_in_box[i][1]
				z=new_pos_in_box[i][2]

				x_new=x-box_x_min
				y_new=y-box_y_min
				z_new=z-box_z_min
				bin_x=int(numpy.floor(x_new/pixel_size))
				bin_y=int(numpy.floor(y_new/pixel_size))
				bin_z=int(numpy.floor(z_new/pixel_size))

				if(bin_x==num_3d_pixel):
					bin_x=num_3d_pixel-1
				if(bin_y==num_3d_pixel): 
					bin_y=num_3d_pixel-1
				if(bin_z==num_3d_pixel):
					bin_z=num_3d_pixel-1      

				if atoms.atom=='N':
					sample[0,bin_x,bin_y,bin_z] = sample[0,bin_x,bin_y,bin_z] + 1
				elif atoms.atom=='CA':
					sample[1,bin_x,bin_y,bin_z] = sample[1,bin_x,bin_y,bin_z] + 1
				elif atoms.atom=='C':
					sample[2,bin_x,bin_y,bin_z] = sample[2,bin_x,bin_y,bin_z] + 1
				elif atoms.atom=='O':
					sample[3,bin_x,bin_y,bin_z] = sample[3,bin_x,bin_y,bin_z] + 1                

			for j in range (0,num_of_channels):
				X_smooth[j,:,:,:]=scipy.ndimage.filters.gaussian_filter(sample[j,:,:,:], sigma=std, order=0, output=None, mode='reflect', cval=0.0, truncate=4.0)
				X_smooth[j,:,:,:]*=1000
			valid = True
	return X_smooth, valid