import sys

from helper import *
from PDB import *
from geom import *
from out_put import *
import top_db

#==main===

if len(sys.argv)<3:
 print 'Usage: python pdb2tcl.py PDB rtfs [ASP] [GLU] [HIS]'
 quit()
#load command line parameters
arg=[]
for i in range(3,len(sys.argv)):
 if sys.argv[i] in ['ASP','GLU','HIS']:
  arg.append(sys.argv[i])
rtf_file=''
for i in range(2,len(sys.argv)):
 if sys.argv[i][-3:] in ['rtf','str','inp']:
  rtf_file+=sys.argv[i]+' '
#====================================
all_system=PDBdata(sys.argv[1])
all_reslist =all_system.reslist
all_system.center()
#decompose the system into components
system_small = PDBdata()
#system_lip = PDBdata()
system=PDBdata()
pro_nm_list=top_db.pro_res_nm_db()
small_nm_list=top_db.sol_res_nm_db()+\
              top_db.lip_res_nm_db()+\
	      top_db.rna_res_nm_db() #temperary solution

sys_list = [[system, pro_nm_list],\
            [system_small, small_nm_list]\
            ]

#print sys_list
seperate_system(all_system, sys_list)

#now for part with small molecules: each has only one residue
#here is assumption is that input PDB has all atom information 
#either from crystal structures or all-atom simulations

#system_small.show_pdb('a.pdb')

if len(system_small.atoms)>0:
 print_title('Solvent, lipid and ion')
 small_system_cg = coarse_graining_system_small(system_small)
 #small_system_cg.show_pdb('a.pdb') 
 if len(small_system_cg.atoms)>0:
  output_build_tcl(rm_directory(sys.argv[1][:-4])+'-small', rtf_file,\
                 small_system_cg, [],'lipids, solvent, ions etc');
  output_build_pdb(rm_directory(sys.argv[1][:-4])+'-small', small_system_cg)

#extract proteon part

if len(system.atoms)==0:
 quit()
print_title('Proteins')
reslist=system.reslist
xp = system.xp

#=========parameters============================================
acceptor_list = system.extract_atompair_byres(top_db.pro_acc_db())

acc_list=[]


donor_list = system.extract_atompair_byres([['N','N'],\
       ['ND2','ND2'],\
       ['NE2','NE2'],\
       ['ND1','ND1'],['NE2','NE2'],['NE1','NE1']])

for i in acceptor_list:
 acc_list.append(i[1])
for i in donor_list:
 acc_list.append(i[1])

check_list= [['HIS',['ND1','NE2'],\
              [['HIS','Nd protonated'],['HISE','Ne protonated'],['HISH','Both protonated']]],\
             ['ASP',['OD1','OD2'],\
              [['ASP','Deprotonated'],['ASPH','Neutral']]],\
             ['GLU',['OE1','OE2'],\
              [['GLU','Deprotonated'],['GLUH','Neutrals']]]]

#=================================================================
#working on protein parts first
#here protein part means a moelcule that has a lot of residues inseatd of a few
#=================================================================
#split chains
seg_st_idx = system.get_segments()
chain_st_idx = system.get_chains()
if seg_st_idx!=chain_st_idx:
 print 'Segment info dosnt match chain info, you may need to resplit chains'
print 'Want to re-split chain structures?(Y/N)'
co = sys.stdin.readline()
if 'Y' in co or 'y' in co:
 reset_chain(system)
 chain_st_idx = system.get_chains()

#reset segname
print 'Want to reset segname?(Y/N)'
co = sys.stdin.readline()
if 'Y' in co or 'y' in co:
 set_segname(system)

#check protonation states
#for i in reslist:
# for ci in check_list:
#  if ci[0] in i[2]:
#   iatom = system.atoms[i[0]]
#   resSig=iatom[8]+':'+iatom[2]+str(iatom[4])
#   print '*************************'
#   print 'In ',
#   print  resSig
#   nm_list=ci[1]
#   xp_list = system.extract_bypart(nm_list,i[0],i[1])
#   for k,l in zip(xp_list, nm_list):
#    print l,'has contacts with',
#    check_donor_dist(system,k,acc_list, resSig, nm_list)
#    print
#   choice=0
#   if ci[0] in arg:
#    print 'Please selected following options for this residue'
#    show_opt_res(ci[2])
#    choice=choose_opt(len(ci[2]))
#   n_resnm= ci[2][choice][0]  
#   for j in range(i[0],i[0]+i[1]):
#    system.atoms[j][2]=n_resnm

#now disulfide bond
print '*************************'
print 'Looking for disulfide bonds...'
cys_S_list = gen_cys_s_list(system)
disu_list =check_disu (system, cys_S_list)

output_build_tcl(rm_directory(sys.argv[1][:-4]), rtf_file, system, disu_list);
output_build_pdb(rm_directory(sys.argv[1][:-4]), system)