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

#==Main==

arg=sys.argv

system = PDBdata(arg[1])
reslist=system.reslist
xp=system.xp

rtf_file=''
for i in range(1,len(arg)):
 if arg[i][-3:] in ['rtf','str','inp']:
  rtf_file+=arg[i]+' '

nname = arg_w_opt(arg,'nname','CL-')
pname = arg_w_opt(arg,'pname','NA+')
conc= float(arg_w_opt(arg,'conc','0.1'))
resrep = arg_w_opt(arg, 'replace', 'SOL')


#===ionize
tc = system.tcharge()
print 'Total charge is ',tc
print 'To neutralize the system and make ',conc,'M salt solution'

vol = system.volume()/1000.
st_M_nm3 = 0.602
n_ion_h = int(conc * st_M_nm3 * vol)/2
nn = n_ion_h
np = n_ion_h
atc =int(abs(tc)/2)
htc =abs(tc)-atc
if tc>0: 
 si=1
elif tc<0:
 si=-1
else:
 si=0

nn+=si*atc
np-=si*htc

print nn,nname,'and',np,pname,'are added'
print 'Residue',resrep,'will be replcaed by ions'
#====
r_list=[]
for i in reslist:
 if i[2]==resrep:
  r_list.append(i[0])

if len(r_list)<nn+np:
 print 'Err! Not enough',resrep,'for replacement'
 quit()

#==select res to replace
random.seed()
rand_list=[]
while len(rand_list)<nn+np:
 ri = random.randint(0,len(r_list)-1)
 if not r_list[ri] in rand_list:
  rand_list.append(r_list[ri])

n_sys = PDBdata()
for i  in reslist:
 if not i[0] in rand_list:
  for j in range(i[0],i[0]+i[1]):
   n_sys.atoms.append(system.atoms[j])

#==gen ions
natom=len(system.atoms)
for i in range(nn):
 r_idx = rand_list[i]
 n_sys.atoms.append([natom+i,nname, nname,'I',i+1,\
                    xp[r_idx][0], xp[r_idx][1], xp[r_idx][2],'ION'])

for i in range(nn, nn+np):
 r_idx = rand_list[i]
 n_sys.atoms.append([natom+i,pname, pname,'I',i+1,\
                    xp[r_idx][0], xp[r_idx][1], xp[r_idx][2],'ION'])

 
n_sys.build_struct()

#==output==

#looking for DISU
cys_S_list = gen_cys_s_list(n_sys)
disu_list =check_disu (n_sys, cys_S_list)

#output
output_build_tcl(rm_directory(arg[1][:-4]), rtf_file,\
                 n_sys, disu_list,'ionized system','_ion');
output_build_pdb(rm_directory(arg[1][:-4]), n_sys)