import numpy
from extract_all_res_defined_AA_type import *
from extract_all_res import *
from cut_box_site_test_pdb import *

from eval_PROSITE_3DCNN import *
from eval_ATP_3DCNN import *
from eval_NOS_3DCNN import *
from eval_TRYPSIN_3DCNN import *

from detection_summary import *
from visualize_detected_sites import * 
import argparse

parser = argparse.ArgumentParser()
parser.add_argument("module", help="require module name")
parser.add_argument("site", help="require site name")
parser.add_argument("pdb_list", help="require pdb_list")

args = parser.parse_args()
module = args.module
site_name = args.site
pdb_list = args.pdb_list

PROSITE_site_info_dict = {
	'PROTEIN_KINASE_ST.5':('ASP','OD2',0.15),
	'PROTEIN_KINASE_TYR.5':('ASP','OD2',0.5),
	'TRYPSIN_HIS.5':('HIS','NE2',0.3),
	'TRYPSIN_SER.6':('SER','OG',0.07),
	'ADH_SHORT.5':('TYR','OH',0.998),
	'EGF_1.10':('CYS','SG',0.9998),
	'IG_MHC.3':('CYS','SG',0.99988),
	'EF_HAND_1.1':('ASP','OD1',0.99),
	'RNASE_PANCREATIC.2':('LYS','NZ',0.07),
	'INSULIN.2':('CYS','SG',0.8)
	}

if module == 'PROSITE':
	total_fold=5
	target_RES, target_ATOM, thres = PROSITE_site_info_dict[site_name]
	ptf_file = extract_all_res_defined_AA_type('../data/pdb/'+pdb_list, '../data/pdb/', target_RES, target_ATOM, site_name)
	numpy_ID = cut_box_site_test_pdb(ptf_file, '../data/pdb/', target_RES, target_ATOM, site_name)
	eval_PROSITE_3DCNN(numpy_ID)
	summary_file = detection_summary(numpy_ID,total_fold)
	visualize_detected_sites([numpy_ID],site_name)

elif module == 'NOS':
	total_fold=5
	all_numpy_ID = []
	for target_RES, target_ATOM in [('ARG','CZ'),('CYS','SG'),('GLU','OE1'),('TRP','NE1')]:
		ptf_file = extract_all_res_defined_AA_type('../data/pdb/'+pdb_list, '../data/pdb/',target_RES, target_ATOM, site_name)
		numpy_ID = cut_box_site_test_pdb(ptf_file, '../data/pdb/', target_RES, target_ATOM, site_name)
		eval_NOS_3DCNN(target_RES, target_ATOM, site_name)
		all_numpy_ID.append(numpy_ID)
		summary_file = detection_summary(numpy_ID,total_fold)
	visualize_detected_sites(all_numpy_ID,site_name)


elif module == 'TRYPSIN_like':
	total_fold=5
	all_numpy_ID = []
	for target_RES, target_ATOM in [('HIS','NE2'),('SER','OG')]:
		ptf_file = extract_all_res_defined_AA_type('../data/pdb/'+pdb_list, '../data/pdb/',target_RES, target_ATOM, site_name)
		numpy_ID = cut_box_site_test_pdb(ptf_file, '../data/pdb/', target_RES, target_ATOM, site_name)
		eval_TRYPSIN_3DCNN(target_RES, target_ATOM, site_name)
		all_numpy_ID.append(numpy_ID)
		summary_file = detection_summary(numpy_ID,total_fold)
	visualize_detected_sites(all_numpy_ID,site_name)

	
elif module == 'ATP':
	total_fold=1
	ptf_file = extract_all_res('../data/pdb/'+pdb_list, '../data/pdb/', site_name)
	numpy_ID = cut_box_site_test_pdb(ptf_file, '../data/pdb/', '', '', site_name)
	eval_ATP_3DCNN(numpy_ID=site_name)
	summary_file = detection_summary(numpy_ID,total_fold)
	visualize_detected_sites([numpy_ID],site_name)